#include "stdio.h"
#include "stdlib.h"
#include "math.h"
#include "mpi.h"

#define a(x,y) a[(x)*M+y]
#define q(x,y) q[(x)*M+y]
#define A(x,y) A[(x)*M+y]
#define Q(x,y) Q[(x)*M+y]
#define R(x,y) R[(x)*M+y]

float temp;
float *A;
float *R;
float *Q;
double starttime;
double time1;
double time2;
int p;
MPI_Status status;

void Environment_Finalize(float *a,float *q,float *v,float *f,float *R,
                          float *Q,float *ai,float *aj,float *qi,float *qj)
{
    free(a);
    free(q);
    free(v);
    free(f);
    free(R);
    free(Q);
    free(ai);
    free(aj);
    free(qi);
    free(qj);
}


int main(int argc, char **argv)
{
    int M,N,m;
    int z,d;
    int i,j,k,my_rank,group_size;
    float *ai,*qi,*aj,*qj;
    float c,s,sp;
    float *f,*v;
    float *a,*q;
    FILE *fdA;

    MPI_Init(&argc,&argv);
    MPI_Comm_rank(MPI_COMM_WORLD,&my_rank);
    MPI_Comm_size(MPI_COMM_WORLD,&group_size);
    p=group_size;
    starttime=MPI_Wtime();
    if(my_rank==p-1)
    {
        fdA=fopen("dataIn.txt","r");
        fscanf(fdA,"%d %d", &M, &N);

        if(M != N)
        {
            puts("The input is error!");
            exit(0);
        }

        A=(float*)malloc(sizeof(float)*M*M);
        Q=(float*)malloc(sizeof(float)*M*M);
        R=(float*)malloc(sizeof(float)*M*M);

        for(i = 0; i < M; i ++)
        {
            for(j = 0; j < M; j ++) fscanf(fdA, "%f", A+i*M+j);
        }
        fclose(fdA);

        for(i=0;i<M;i++)
            for(j=0;j<M;j++)
                if (i==j)
                    Q(i,j)=1.0;
        else
            Q(i,j)=0.0;
		printf("rank %d got the matrix:\n", my_rank);
		for(i=0;i<M;i++){
			for(j=0;j<M;j++)
				printf("%f ",A(i,j));
			printf("\n");
		}
    }
    MPI_Bcast(&M,1,MPI_INT,p-1,MPI_COMM_WORLD);
    m=M/p;
    if (M%p!=0) m++;

    qi=(float*)malloc(sizeof(float)*M);
    qj=(float*)malloc(sizeof(float)*M);
    aj=(float*)malloc(sizeof(float)*M);
    ai=(float*)malloc(sizeof(float)*M);
    v=(float*)malloc(sizeof(float)*M);
    f=(float*)malloc(sizeof(float)*M);
    a=(float*)malloc(sizeof(float)*m*M);
    q=(float*)malloc(sizeof(float)*m*M);

    if (a==NULL||q==NULL||f==NULL||v==NULL||qi==NULL||qj==NULL||ai==NULL||aj==NULL)
        printf("memory allocation is wrong\n");


	//隔行分發
	if(my_rank==p-1){ 
		for(i=0;i<M;i++){
			MPI_Send(&A(i,0),M,MPI_FLOAT,i%p,i/p,MPI_COMM_WORLD);
			MPI_Send(&Q(i,0),M,MPI_FLOAT,i%p,i/p,MPI_COMM_WORLD);
		}
	}
	for(i=0;i<m;i++){
		MPI_Recv(&a(i,0),M,MPI_FLOAT,p-1,i,MPI_COMM_WORLD,&status);
		MPI_Recv(&q(i,0),M,MPI_FLOAT,p-1,i,MPI_COMM_WORLD,&status);
	}

	time1=MPI_Wtime();


	if(p>1)
	{
		for(i=0;i<m;i++)
		{
			for(j=0;j<p;j++)
			{
				if(my_rank==j)
				{
					for(k=0;k<M;k++)
					{
						f[k]=a(i,k);
						v[k]=q(i,k);
					}		
					MPI_Send(&f[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
					MPI_Send(&v[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
				}
				if(my_rank<=j)
				{
					for(d=0;d<m-i-1;d++)
					{
						MPI_Recv(&f[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
						MPI_Recv(&v[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
						sp=sqrt(f[i*p+j]*f[i*p+j]+a(i+d+1,i*p+j)*a(i+d+1,i*p+j));
						c=f[i*p+j]/sp;s=a(i+d+1,i*p+j)/sp;
						for(k=0;k<M;k++)
						{
							aj[k]=c*f[k]+s*a(i+d+1,k);
							qj[k]=c*v[k]+s*q(i+d+1,k);
							ai[k]=-s*f[k]+c*a(i+d+1,k);
							qi[k]=-s*v[k]+c*q(i+d+1,k);
						}
						for(k=0;k<M;k++)
						{
							f[k]=aj[k];
							v[k]=qj[k];
							a(i+d+1,k)=ai[k];
							q(i+d+1,k)=qi[k];
						}
						if((((M%p-1)!=0)&&my_rank!=(M%p-1))||(my_rank==(M%p-1)&&d!=m-i-2)){
							MPI_Send(&f[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
							MPI_Send(&v[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
						}
					}
				}
				if(my_rank>j)
				{
					for(d=0;d<m-i;d++)
					{
						MPI_Recv(&f[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
						MPI_Recv(&v[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
						sp=sqrt(f[i*p+j]*f[i*p+j]+a(i+d,i*p+j)*a(i+d,i*p+j));
						c=f[i*p+j]/sp;s=a(i+d,i*p+j)/sp;
						for(k=0;k<M;k++)
						{
							aj[k]=c*f[k]+s*a(i+d,k);
							qj[k]=c*v[k]+s*q(i+d,k);
							ai[k]=-s*f[k]+c*a(i+d,k);
							qi[k]=-s*v[k]+c*q(i+d,k);
						}
						for(k=0;k<M;k++)
						{
							f[k]=aj[k];
							v[k]=qj[k];
							a(i+d,k)=ai[k];
							q(i+d,k)=qi[k];
						}
						if((M%p==0&&my_rank!=(M%p-1))||(my_rank==(M%p-1)&&d!=m-i-1)){
							MPI_Send(&f[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
							MPI_Send(&v[0],M,MPI_FLOAT,(my_rank+1)%p,i,MPI_COMM_WORLD);
						}
					}				
				}
				if(my_rank==M%p-1)
				{
					MPI_Send(&f[0],M,MPI_FLOAT,p-1,i,MPI_COMM_WORLD);
					MPI_Send(&v[0],M,MPI_FLOAT,p-1,i,MPI_COMM_WORLD);
				}
				if(my_rank==p-1)
				{
					if(M%p-1!=0)
					{
						MPI_Recv(&f[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
						MPI_Recv(&v[0],M,MPI_FLOAT,(my_rank-1)%p,i,MPI_COMM_WORLD,&status);
					}
                    for(k=0;k<M;k++)
                    {
                        Q(i,k)=v[k];
                        R(i,k)=f[k];
                    }				
				}
			}
		}
	}
    if (p==1)
    {
        for (j=0;j<M;j++)
            for (i=j+1;i<M;i++)
        {
            sp=sqrt(a(j,j)*a(j,j) + a(i,j)*a(i,j));
            c=a(j,j)/sp;
            s=a(i,j)/sp;

            for (k=0;k<M;k++)
            {
                aj[k]=c*a(j,k) + s*a(i,k);
                qj[k]=c*q(j,k) + s*q(i,k);
                ai[k]=(-s)*a(j,k) + c*a(i,k);
                qi[k]=(-s)*q(j,k) + c*q(i,k);
            }

            for (k=0;k<M;k++)
            {
                a(j,k)=aj[k];
                q(j,k)=qj[k];
                a(i,k)=ai[k];
                q(i,k)=qi[k];
            }
        }                                

        for(i=0;i<M;i++)
            for(j=0;j<M;j++)
                R(i,j)=a(i,j);

        for(i=0;i<M;i++)
            for(j=0;j<M;j++)
                Q(i,j)=q(i,j);
    }                         

    if (my_rank==p-1)
    {
        printf("Input of file \"dataIn.txt\"\n");
        printf("%d\t %d\n",M, N);
        for(i=0;i<M;i++)
        {
            for(j=0;j<N;j++) printf("%f\t",A(i,j));
            printf("\n");
        }

        printf("\nOutput of QR operation\n");

        printf("Matrix R:\n");
        for(i=0;i<M;i++)
        {
            for(j=0;j<M;j++)
                printf("%f\t",R(i,j));
            printf("\n");
        }

        for(i=0;i<M;i++)
            for(j=i+1;j<M;j++)
        {
            temp=Q(i,j);
            Q(i,j)=Q(j,i);
            Q(j,i)=temp;
        }

        printf("Matrix Q:\n");
        for(i=0;i<M;i++)
        {
            for(j=0;j<M;j++)
                printf("%f\t",Q(i,j));
            printf("\n");
        }
    }

    time2 = MPI_Wtime();
    if (my_rank==0)
    {
        printf("\n");
        printf("Whole running time    = %f seconds\n",time2-starttime);
        printf("Distribute data time  = %f seconds\n",time1-starttime);
        printf("Parallel compute time = %f seconds\n",time2-time1);

    }

    MPI_Barrier(MPI_COMM_WORLD);
    MPI_Finalize();
    Environment_Finalize(a,q,v,f,R,Q,ai,aj,qi,qj);
    return(0);
}