/*
************************************************************************
                        Example 25 (congrad.c)

       Objective            : Conjugate Gradient Method to solve AX = b
                              matrix system of linear equations.

       Input                : Real Symmetric Positive definite Matrix
                              and the real vector - Input_A and Vector_B
                              Read files (mdatcg.inp) for Matrix A
                              and (vdatcg.inp) for Vector b

       Description          : Input matrix is stored in n by n format.
                              Diagonal preconditioning matrix is used.
                              Rowwise block striped partitioning matrix
                              is used.Maximum iterations is given by 
                              MAX_ITERATIONS.Tolerance value is given 
                              by EPSILON
                              Header file used  : cg_constants.h

       Output               : The solution of  Ax=b on process with 
                              Rank 0 and the number of iterations 
                              for convergence of the method.

       Necessary conditions : Number of Processes should be less than
                              or equal to 8 Input Matrix Should be 
                              Square Matrix. Matrix size for Matrix A
                              and vector size for vector b should be 
                              equally striped, that is Matrix size and
                              Vector Size should be dividible by the 
                              number of processes used.

***************************************************************
*/



#include<math.h>
#include<stdio.h>
#include<stdlib.h>
#include<mpi.h>

#define EPSILON 		  1.0E-20
#define MAX_ITERATIONS 10000

/******************************************************************************/
void 
GetPreconditionMatrix(double **Bloc_Precond_Matrix, int NoofRows_Bloc, 
		      int NoofCols)
{

	/*... Preconditional Martix is identity matrix .......*/
	int 	Bloc_MatrixSize;
	int 	irow, icol, index;
	double *Precond_Matrix;

	Bloc_MatrixSize = NoofRows_Bloc*NoofCols;

	Precond_Matrix = (double *) malloc(Bloc_MatrixSize * sizeof(double));

   index = 0;
	for(irow=0; irow<NoofRows_Bloc; irow++){
		for(icol=0; icol<NoofCols; icol++){
			Precond_Matrix[index++] = 1.0;
		}
	}
	*Bloc_Precond_Matrix = Precond_Matrix;
}
/******************************************************************************/
double 
ComputeVectorDotProduct(double *Vector1, double *Vector2, int VectorSize)
{
	int 	index;
	double Product;

	Product = 0.0;
	for(index=0; index<VectorSize; index++)
		Product += Vector1[index]*Vector2[index];

	return(Product);
}
/******************************************************************************/
void 
CalculateResidueVector(double *Bloc_Residue_Vector, double *Bloc_Matrix_A, 
							  double *Input_B, double *Vector_X, int NoofRows_Bloc, 
							  int VectorSize, int MyRank)
{
	/*... Computes residue = AX - b .......*/
	int   irow, index, GlobalVectorIndex;
	double value;

	GlobalVectorIndex = MyRank * NoofRows_Bloc;
	for(irow=0; irow<NoofRows_Bloc; irow++){
		index = irow * VectorSize;
		value = ComputeVectorDotProduct(&Bloc_Matrix_A[index], Vector_X, 
												  VectorSize);
		Bloc_Residue_Vector[irow] = value - Input_B[GlobalVectorIndex++];
	}
}
/******************************************************************************/
void
SolvePrecondMatrix(double *Bloc_Precond_Matrix, double *HVector, 
						 double *Bloc_Residue_Vector, int Bloc_VectorSize)
{
	/*...HVector = Bloc_Precond_Matrix inverse * Bloc_Residue_Vector.......*/
	int	index;

	for(index=0; index<Bloc_VectorSize; index++){
		HVector[index] = Bloc_Residue_Vector[index]/1.0;
	}
}
/******************************************************************************/
main(int argc, char *argv[])
{
	int        NumProcs, MyRank, Root=0;
	int        NoofRows, NoofCols, VectorSize;
	int        n_size, NoofRows_Bloc, Bloc_MatrixSize, Bloc_VectorSize;
	int	     Iteration = 0, irow, icol, index, CorrectResult;
	double	   **Matrix_A, *Input_A, *Input_B, *Vector_X, *Bloc_Vector_X; 
	double      *Bloc_Matrix_A, *Bloc_Precond_Matrix, *Buffer;
	double 	  *Bloc_Residue_Vector ,*Bloc_HVector, *Bloc_Gradient_Vector;
	double		  *Direction_Vector, *Bloc_Direction_Vector;
	double      Delta0, Delta1, Bloc_Delta0, Bloc_Delta1;
	double		  Tau, val, temp, Beta;

	double      *AMatDirect_local, *XVector_local;
	double      *ResidueVector_local, *DirectionVector_local;
	double     StartTime, EndTime;
	MPI_Status status;
	FILE  	  *fp;

  	/*...Initialising MPI .......*/
	MPI_Init(&argc,&argv);
	MPI_Comm_size(MPI_COMM_WORLD,&NumProcs);
	MPI_Comm_rank(MPI_COMM_WORLD,&MyRank);


  /* .......Read the Input file ......*/
  if(MyRank == Root) {

     if ((fp = fopen ("./data/mdatcg.inp", "r")) == NULL) {
       printf("Can't open input matrix file");
       exit(-1);
     }
     fscanf(fp, "%d %d", &NoofRows,&NoofCols);     
     n_size=NoofRows;

     /* ...Allocate memory and read data .....*/
	  Matrix_A = (double **) malloc(n_size*sizeof(double *));

     for(irow = 0; irow < n_size; irow++){
		  Matrix_A[irow] = (double *) malloc(n_size * sizeof(double));
        for(icol = 0; icol < n_size; icol++)
	        fscanf(fp, "%lf", &Matrix_A[irow][icol]);
	  }
     fclose(fp);

     if ((fp = fopen ("./data/vdatcg.inp", "r")) == NULL){
        printf("Can't open input vector file");
        exit(-1);
     }

     fscanf(fp, "%d", &VectorSize);     
     n_size=VectorSize;
     Input_B  = (double *)malloc(n_size*sizeof(double));
     for (irow = 0; irow<n_size; irow++)
        fscanf(fp, "%lf",&Input_B[irow]);
     fclose(fp); 

	  /* ...Convert Matrix_A into 1-D array Input_A ......*/
     Input_A  = (double *)malloc(n_size*n_size*sizeof(double));
	  index    = 0;
	  for(irow=0; irow<n_size; irow++)
	  	  for(icol=0; icol<n_size; icol++)
			  Input_A[index++] = Matrix_A[irow][icol];

  }

	MPI_Barrier(MPI_COMM_WORLD);
	StartTime = MPI_Wtime();

	/*...Broadcast Matrix and Vector size and perform input validation tests...*/
	MPI_Bcast(&NoofRows, 1, MPI_INT, Root, MPI_COMM_WORLD);
	MPI_Bcast(&NoofCols, 1, MPI_INT, Root, MPI_COMM_WORLD);
	MPI_Bcast(&VectorSize, 1, MPI_INT, Root, MPI_COMM_WORLD);


   if(NoofRows != NoofCols){
		MPI_Finalize();
		if(MyRank == Root)
			printf("Error : Coefficient Matrix Should be square matrix");
		exit(-1);
	}

   if(NoofRows != VectorSize){
		MPI_Finalize();
		if(MyRank == Root)
			printf("Error : Matrix Size should be equal to VectorSize");
		exit(-1);
	}

	if(NoofRows % NumProcs != 0){
		MPI_Finalize();
		if(MyRank == Root)
			printf("Error : Matrix cannot be evenly striped among processes");
		exit(-1);
	}

   /*...Allocate memory for Input_B and BroadCast Input_B.......*/
   if(MyRank != Root)
		Input_B = (double *) malloc(VectorSize*sizeof(double));
	MPI_Bcast(Input_B, VectorSize, MPI_DOUBLE, Root, MPI_COMM_WORLD);

   /*...Allocate memory for Block Matrix A and Scatter Input_A .......*/
   NoofRows_Bloc   = NoofRows / NumProcs;
	Bloc_VectorSize = NoofRows_Bloc;
	Bloc_MatrixSize = NoofRows_Bloc * NoofCols;
	Bloc_Matrix_A = (double *) malloc (Bloc_MatrixSize*sizeof(double));
	MPI_Scatter(Input_A, Bloc_MatrixSize, MPI_DOUBLE,
				   Bloc_Matrix_A, Bloc_MatrixSize, MPI_DOUBLE, Root, MPI_COMM_WORLD);


	/*... Allocates memory for solution vector and intialise it to zero.......*/
	Vector_X = (double *) malloc(VectorSize*sizeof(double));
	for(index=0; index<VectorSize; index++)
		Vector_X[index] = 0.0;

	/*...Calculate RESIDUE = AX - b .......*/
   Bloc_Residue_Vector = (double *) malloc(NoofRows_Bloc*sizeof(double));
	CalculateResidueVector(Bloc_Residue_Vector,Bloc_Matrix_A, Input_B, Vector_X,
								  NoofRows_Bloc, VectorSize, MyRank);

	/*... Precondtion Matrix is identity matrix ......*/
	GetPreconditionMatrix( &Bloc_Precond_Matrix, NoofRows_Bloc, NoofCols);

	/*...Bloc_HVector = Bloc_Precond_Matrix inverse * Bloc_Residue_Vector......*/
	Bloc_HVector = (double *) malloc(Bloc_VectorSize*sizeof(double));
	SolvePrecondMatrix(Bloc_Precond_Matrix, Bloc_HVector, Bloc_Residue_Vector, 
							 Bloc_VectorSize); 


   /*...Initailise Bloc Direction Vector = -(Bloc_HVector).......*/	
	Bloc_Direction_Vector = (double *) malloc(Bloc_VectorSize*sizeof(double));
	for(index=0; index<Bloc_VectorSize; index++)
		Bloc_Direction_Vector[index] = 0 - Bloc_HVector[index];

	/*...Calculate Delta0 and check for convergence .......*/
	Bloc_Delta0 = ComputeVectorDotProduct(Bloc_Residue_Vector,
										Bloc_HVector, Bloc_VectorSize);
	MPI_Allreduce(&Bloc_Delta0, &Delta0, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

	if(Delta0 < EPSILON){
		MPI_Finalize();
		exit(0);
	}

	/*...Allocate memory for Direction Vector.......*/
	Direction_Vector = (double *) malloc(VectorSize*sizeof(double));

	/*...Allocate temporary buffer to store Bloc_Matrix_A*Direction_Vector...*/
	Buffer = (double *) malloc(Bloc_VectorSize*sizeof(double));

	/*...Allocate memory for Bloc_Vector_X .......*/
	Bloc_Vector_X = (double *)malloc(Bloc_VectorSize*sizeof(double));

	Iteration = 0;
	do{

		Iteration++;
		/*
		if(MyRank == Root)
			printf("Iteration : %d\n",Iteration);
		*/

		/*...Gather Direction Vector on all processes.......*/
		MPI_Allgather(Bloc_Direction_Vector, Bloc_VectorSize, MPI_DOUBLE, 
		            Direction_Vector, Bloc_VectorSize, MPI_DOUBLE, MPI_COMM_WORLD);

		/*...Compute Tau = Delta0 / (DirVector Transpose*Matrix_A*DirVector)...*/
		for(index=0; index<NoofRows_Bloc; index++){
			Buffer[index] = ComputeVectorDotProduct(&Bloc_Matrix_A[index*NoofCols],
													  Direction_Vector, VectorSize);
		}
		temp = ComputeVectorDotProduct(Bloc_Direction_Vector, Buffer, 
												 Bloc_VectorSize);

		MPI_Allreduce(&temp, &val, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
		Tau = Delta0 / val;

		/*Compute new vector Xnew = Xold + Tau*Direction........................*/
		/*Compute BlocResidueVec  = BlocResidueVect + Tau*Bloc_MatA*DirVector...*/
		for(index = 0; index<Bloc_VectorSize; index++){
			Bloc_Vector_X[index]       = Vector_X[MyRank*Bloc_VectorSize + index] +
										  		  Tau*Bloc_Direction_Vector[index];
			Bloc_Residue_Vector[index] = Bloc_Residue_Vector[index] + 
												  Tau*Buffer[index];
		}
		/*...Gather New Vector X at all  processes......*/
		MPI_Allgather(Bloc_Vector_X, Bloc_VectorSize, MPI_DOUBLE, Vector_X, 
						  Bloc_VectorSize, MPI_DOUBLE, MPI_COMM_WORLD);

		SolvePrecondMatrix(Bloc_Precond_Matrix, Bloc_HVector, Bloc_Residue_Vector,
								 Bloc_VectorSize); 
		Bloc_Delta1 = ComputeVectorDotProduct(Bloc_Residue_Vector, Bloc_HVector, 
														  Bloc_VectorSize);
		
		MPI_Allreduce(&Bloc_Delta1, &Delta1, 1, MPI_DOUBLE, MPI_SUM, 
						  MPI_COMM_WORLD);

		if(Delta1 < EPSILON)
			break;

		Beta   = Delta1 / Delta0;
		Delta0 = Delta1;
		for(index=0; index<Bloc_VectorSize; index++){
			Bloc_Direction_Vector[index] = -Bloc_HVector[index] + 
													 Beta*Bloc_Direction_Vector[index];
		}
	}while(Delta0 > EPSILON && Iteration < MAX_ITERATIONS);

	MPI_Barrier(MPI_COMM_WORLD);
	EndTime = MPI_Wtime();

  if (MyRank == 0) {

     printf ("\n");
     printf(" ------------------------------------------- \n");
     printf("Results of Jacobi Method on processor %d: \n", MyRank);
     printf ("\n");

     printf("Matrix Input_A \n");
     printf ("\n");
     for (irow = 0; irow < n_size; irow++) {
        for (icol = 0; icol < n_size; icol++)
	        printf("%.3lf  ", Matrix_A[irow][icol]);
        printf("\n");
     }
     printf ("\n");
     printf("Matrix Input_B \n");
     printf("\n");
     for (irow = 0; irow < n_size; irow++) {
         printf("%.3lf\n", Input_B[irow]);
     }
     printf ("\n");
     printf("Solution vector \n");
	  printf("Number of iterations = %d\n",Iteration);
     printf ("\n");
     for(irow = 0; irow < n_size; irow++)
        printf("%.12lf\n",Vector_X[irow]);
     printf(" --------------------------------------------------- \n");
  }
  


	MPI_Finalize();
}