// ガウスの消去法で連立一次方程式を解くプログラム

// 4x + y + z = 9
// x + 3y + z = 10
// 2x + y + 5z = 19

#include <stdio.h>
#include "mpi.h"

#define SIZE 3

double A[SIZE][SIZE] = {
	{4, 1, 1},
	{1, 3, 1},
	{2, 1, 5}
};

double b[SIZE]={9,10,19};

double c[SIZE];
double x[SIZE];

int debug_id = 0;


int main(int argc, char* argv[]){
	int i,j,k;
	int n = SIZE;

	double L[SIZE];
	
	MPI_Status istatus;
	int myid, numprocs;

	MPI_Init(&argc, &argv);
	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
	MPI_Comm_size(MPI_COMM_WORLD, &numprocs);

	// 1つのプロセスが担当する列の長さ
	int partition_size = n / numprocs;
	// プロセスが担当する列の開始点
	int start = myid * partition_size;
	// プロセスが担当する列の終了点
	int end = (myid + 1) * partition_size;

	// 外積方式ガウス法

	// LU分解
	for (k=0; k<n; k++) {
		// kで回す行方向ループ

		int worker = k / partition_size;

		// 行列Lの処理
		if(worker == myid){
			// kが自分の担当範囲内の場合はLを処理
			double div_tmp = 1.0 / A[k][k];
			for (i=k+1; i<n; i++) {
				// iで回す行方向ループ

				// l_ik = a_ik / a_kk
				A[i][k] = A[i][k] * div_tmp;
				L[i] = A[i][k];
			}
		}


		// Lを放送
		MPI_Bcast(&L[k], n - k, MPI_DOUBLE, worker, MPI_COMM_WORLD);


		// 行列Uの処理
		for (j=k+1; j<n; j++) {
			// jで回す行方向のループ

			double a_jk = L[j];

			// 並列処理
			if(k < start){
				i = start;
			}
			else{
				i = k+1;
			}
			
			for (; i < end; i++) {
				// iで回す列方向のループ
				A[j][i] = A[j][i] - A[k][i] * a_jk;
			}
		}

		MPI_Barrier(MPI_COMM_WORLD);
	}
	// LU分解終了


	// 結果を集める
	for(i = 0; i < n; i++){
		// 行方向のループ
		for(j = partition_size; j < n; j+=partition_size){
			// 列方向のループ

			int worker = j / partition_size;

			if(myid == 0){
				// process0 はその他のプロセスの担当部分を集める
				MPI_Recv(&A[i][j], partition_size, MPI_DOUBLE, worker, 0, MPI_COMM_WORLD, &istatus);
			}
			if(myid == worker){
				// process0 以外は、自分が担当した部分を送る
				MPI_Send(&A[i][j], partition_size, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
			}
			MPI_Barrier(MPI_COMM_WORLD);
		}
	}

	// LU分解終了

	MPI_Finalize();


	// LU分解後のLUを出力
	if(myid == 0){
		printf("LU:\n");
		for(i=0;i<n;i++){
			for(j=0;j<n;j++){
				printf("% 2.2f ",A[i][j]);
			}
			printf("\n");
		}
		printf("\n");
	}


	// 前進代入
	// Lc = b
	for(k=0; k<n; k++){
		// 行方向のループ
		c[k] = b[k];
		for(i=0; i<k; i++){
			// 列方向のループ
			c[k] -= A[k][i] * c[i];
		}
	}


	// 後進代入
	// Ux = c
	// k = n - 1
	x[n-1] = c[n-1] / A[n-1][n-1];
	// k = (n - 2) ~ (0)
	for(k = n - 2; k >= 0; k--){
		x[k] = c[k];
		for(i=k+1; i<n; i++){
			// 列方向のループ
			x[k] -= A[k][i] * x[i];
		}
		x[k] = x[k] / A[k][k];
	}

	if(myid == 0){
		// xを出力
		printf("x:\n");
		for(i=0;i<n;i++){
			printf("% 2.2f ",x[i]);
		}
		printf("\n");
	}


	return 0;
}