#include "mat_mul.h"

#include <cstdio>
#include <cstdlib>
#include <mpi.h>
#include <omp.h>
#include "util.h"

#define min(a, b) (((a) < (b)) ? (a) : (b))
#define MASTER 0
#define FORM_MASTER 1
#define FORM_WORKER 2


static float *A, *B, *C;
static int M, N, K;
static int num_threads;
static int mpi_rank, mpi_world_size;
static int rows[4]={0,};
static int offset[4]={0,};

#define SLICEM 32
#define SLICEK 16
#define SLICEN 2048

static void mat_mul_omp() {
  // TODO: parallelize & optimize matrix multiplication
  // Use num_threads per node

// #pragma omp parallel for
//   for (int i = 0; i < rows; ++i) {
//     for (int j = 0; j < N; ++j) {
//       for (int k = 0; k < K; ++k) {
//         C[i * N + j] += A[i * K + k] * B[k * N + j];
//       }
//     }
//   }

// #pragma omp parallel num_threads(num_threads)
// #pragma omp parallel for
//   for (int i = 0; i < rows; ++i) {
//     for (int k = 0; k < K; ++k) {
//       float arr = A[i*K+k];
//       for (int j = 0; j < N; j+=1) {
//         C[i * N + j] += arr * B[k * N + j];
//       }
//     }
//   }
  int start=0;
  int ed=rows[mpi_rank];

  #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
    for(int i2=start;i2<ed;i2+=SLICEM) {
      for(int k2=0;k2<K;k2+=SLICEK) {
        for(int j2=0;j2<N;j2+=SLICEN) {
          int edk=k2+SLICEK<K?(k2+SLICEK):K;
          int edm=i2+SLICEM<M?(i2+SLICEM):M;
          int edn=j2+SLICEN<N?(j2+SLICEN):N;
          for(int i=i2;i<edm;++i) {
            for(int k=k2;k<edk;++k) {
              for(int j=j2;j<edn;++j) {
                C[i*N+j]+=A[i*K+k]*B[k*N+j];
              }
            }
          }
        }
      }
    }
  }
  // omp_set_num_threads(num_threads);
  // #pragma omp parallel 
  // {
  // int pid = omp_get_thread_num();
  // //printf("Parallel section: Hello world from thread %d\n",pid);

  // int slice = M / num_threads;
  // int start = pid * slice;
  // int end = pid == num_threads - 1 ? M : (pid + 1) * slice;
  // //int unroll = 8;
  // float Aik;
  // int bs = 45;

  // for (int kk = 0; kk < K; kk += bs) {
  //   //for (int jj = 0; jj < N; jj += bs) {
  //     for (int i = start; i < end; ++i) {
  //       for (int k = kk; k < min(kk + bs, K); ++k) {
  //         Aik = A[i * K + k];
  //         int j=0;

  //         for(; j < N; ++j) {
  //           C[i * N + j] += Aik * B[k * N + j];
  //         }
  //       }
  //     }
  //   }
  // }


void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K,
             int _num_threads, int _mpi_rank, int _mpi_world_size) {
  A = _A, B = _B, C = _C;
  M = _M, N = _N, K = _K;
  num_threads = _num_threads, mpi_rank = _mpi_rank,
  mpi_world_size = _mpi_world_size;


  // TODO: parallelize & optimize matrix multiplication on multi-node
  // You must allocate & initialize A, B, C for non-root processes
  //int numworkers = mpi_world_size -1;
  //int averow, extra;
  MPI_Status status;
  MPI_Request request;
  // FIXME: for now, only root process runs the matrix multiplication.
  // if (mpi_rank == 0)
  //   mat_mul_omp();
  int nrows = M / mpi_world_size;
  //averow = M/numworkers;
  //extra = M%numworkers;
  //offset = 0;
   
   for(int i=0;i<mpi_world_size;i++) {
     rows[i]=(i==mpi_world_size-1)?(M-(nrows *(mpi_world_size-1))):nrows;
   }
   for(int i=0;i<mpi_world_size-1;i++) {
     offset[i+1]=offset[i]+rows[i];
   }
   if(mpi_rank != 0) {
     M=rows[mpi_rank];
     alloc_mat(&A, rows[mpi_rank], K);
     alloc_mat(&B, K, N);
     alloc_mat(&C, rows[mpi_rank], N);     
   }
   MPI_Bcast(B,K*N,MPI_FLOAT,0,MPI_COMM_WORLD);
   if(mpi_rank == 0) {
     for(int i=1;i<mpi_world_size;i++)
       MPI_Isend(&A[offset[i]*K],rows[i]*K,MPI_FLOAT,i,0,MPI_COMM_WORLD,&request);
   } else {
       MPI_Recv(A,rows[mpi_rank]*K,MPI_FLOAT,0,0,MPI_COMM_WORLD,&status);
   }
   mat_mul_omp();
   if(mpi_rank != 0) {
       MPI_Isend(C,rows[mpi_rank]*N,MPI_FLOAT,0,0,MPI_COMM_WORLD,&request);
   } else {
     for(int i=1;i<mpi_world_size;i++)
       MPI_Recv(&C[offset[i]*N],rows[i]*N,MPI_FLOAT,i,0,MPI_COMM_WORLD,&status);
   }


}

  //  }     
  //  )
  //  {
  //    /* code */
  //  }
   
  //   printf("/////////////////numworkers %d\n",numworkers);
  //   printf("/////////////////averow %d\n",averow);
  //   printf("/////////////////extra %d\n",extra);

  // if(mpi_rank == 0)
  // {
  //   int row_size = M / mpi_world_size;
  //   int st,ed;
  //   for (int node=1; node < mpi_world_size; node++) {
  //     st = offset = node * row_size;
  //     ed = node == mpi_world_size -1 ? M : (node+1)*row_size;
  //     rows = ed - st;
	// 	MPI_Isend(&offset, 1, MPI_INT, node, FORM_MASTER, MPI_COMM_WORLD, &request);
	// 	MPI_Isend(&rows, 1, MPI_INT, node, FORM_MASTER, MPI_COMM_WORLD,&request);
	// 	MPI_Isend(&A[offset*K], rows*K, MPI_FLOAT, node, FORM_MASTER, MPI_COMM_WORLD,&request);
	// 	MPI_Isend(B, N*K, MPI_FLOAT, node, FORM_MASTER, MPI_COMM_WORLD,&request);

  //   }
  //   rows = row_size;
  //   mat_mul_omp();
  //   for (int node=1; node < mpi_world_size; node++) {
	// 	MPI_Recv(&offset, 1, MPI_INT, node, FORM_WORKER, MPI_COMM_WORLD, &status);
	// 	MPI_Recv(&rows, 1, MPI_INT, node, FORM_WORKER, MPI_COMM_WORLD, &status);
	// 	MPI_Recv(&C[offset*N], rows*N, MPI_FLOAT, node, FORM_WORKER, MPI_COMM_WORLD, &status);

  //   }
  // } else {
  //   alloc_mat(&A, M, K);
  //   alloc_mat(&B, M, N);
  //   alloc_mat(&C, M, N);
  //   zero_mat(C,M,N);

	// 	MPI_Recv(&offset, 1, MPI_INT, MASTER, FORM_MASTER, MPI_COMM_WORLD, &status);
	// 	MPI_Recv(&rows, 1, MPI_INT, MASTER, FORM_MASTER, MPI_COMM_WORLD, &status);
	// 	MPI_Recv(&A, rows*K, MPI_FLOAT, MASTER, FORM_MASTER, MPI_COMM_WORLD, &status);
	// 	MPI_Recv(B, N*K, MPI_FLOAT, MASTER, FORM_MASTER, MPI_COMM_WORLD, &status);
    
  //   mat_mul_omp();
	//   MPI_Isend(&offset, 1, MPI_INT, MASTER, FORM_WORKER, MPI_COMM_WORLD,&request);
  //   MPI_Isend(&rows, 1, MPI_INT, MASTER, FORM_WORKER, MPI_COMM_WORLD,&request);
  //   MPI_Isend(&C, rows*N, MPI_FLOAT, MASTER, FORM_WORKER, MPI_COMM_WORLD,&request);
 
  // }

//     for (dest=1; dest<=numworkers; dest++)
// 	  {
// 		rows = (dest <= extra) ? averow+1 : averow;
// 		// MPI_Send(&offset, 1, MPI_INT, dest, mtype, MPI_COMM_WORLD);
// 		// MPI_Send(&rows, 1, MPI_INT, dest, mtype, MPI_COMM_WORLD);
// 		// MPI_Send(&A[offset], rows*N, MPI_FLOAT, dest, mtype, MPI_COMM_WORLD);
// 		// MPI_Send(B, N*K, MPI_FLOAT, dest, mtype, MPI_COMM_WORLD);
// 		offset = offset + rows;
//     printf(" \n");
//     printf("///////////////// mpi_rank Number %d\n",mpi_rank);
//     printf("///////////////// Master numworkers %d\n",numworkers);
//     printf("///////////////// Master dest %d\n",dest);
//     printf("///////////////// Master send offset %d\n",offset);
//     printf("///////////////// Master send rows %d\n",rows);
//     printf("///////////////// Master send rows*N %d\n",rows*N);
//     printf("///////////////// Master send N*K %d\n",N*K);
//     printf("///////////////// ------------------------------\n");
// 	 }

//   mat_mul_omp();
// 	mtype = FORM_WORKER;
// 	for(i=1; i<=numworkers; i++)
// 	{
// 		source = i;
// 		// MPI_Recv(&offset, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);
// 		// MPI_Recv(&rows, 1, MPI_INT, source, mtype, MPI_COMM_WORLD, &status);
// 		// MPI_Recv(&C[offset], rows*N, MPI_FLOAT, source, mtype, MPI_COMM_WORLD, &status);
//     printf("///////////////// mpi_rank Number %d\n",mpi_rank);
//     printf("///////////////// Master source  %d\n",source);
//     printf("///////////////// Master Recv offset %d\n",offset);
//     printf("///////////////// Master Recv rows %d\n",rows);
//     printf("///////////////// Master Recv C rows*N %d\n",rows*N);
//     printf("///////////////// ------------------------------\n");
// 	}
// }
//   // if (mpi_rank > MASTER)
//   // {
// 	//   mtype = FORM_MASTER;
// 	// 	// MPI_Recv(&offset, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD, &status);
// 	// 	// MPI_Recv(&rows, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD, &status);
// 	// 	// MPI_Recv(&A, rows*N, MPI_DOUBLE, MASTER, mtype, MPI_COMM_WORLD, &status);
// 	// 	// MPI_Recv(B, N*K, MPI_DOUBLE, MASTER, mtype, MPI_COMM_WORLD, &status);
//   //   printf("///////////////// mpi_rank Number %d\n",mpi_rank);
//   //   printf("///////////////// Worker Recv offset %d\n",offset);
//   //   printf("///////////////// Worker Recv rows %d\n",rows);
//   //   printf("///////////////// Worker Recv A rows*N %d\n",rows*N);
//   //   printf("///////////////// Worker Recv B N*K %d\n",N*K);
//   //   printf("///////////////// ------------------------------\n");

// 	//     mat_mul_omp();

// 	//   mtype = FORM_WORKER;
// 	//   // MPI_Send(&offset, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD);
//   //   // MPI_Send(&rows, 1, MPI_INT, MASTER, mtype, MPI_COMM_WORLD);
//   //   // MPI_Send(&C, N, MPI_DOUBLE, MASTER, mtype, MPI_COMM_WORLD);
//   //   printf("///////////////// mpi_rank Number %d\n",mpi_rank);
//   //   printf("///////////////// Worker Send offset %d\n",offset);
//   //   printf("///////////////// Worker Send rows %d\n",rows);
//   //   printf("///////////////// Worker Send C N  %d\n",N);
//   //   printf("///////////////// ------------------------------\n");


//   // }