#include "mat_mul.h"

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

#include "util.h"
#define MASTER 0
#define FROM_MASTER 1
#define FROM_WORKER 2

static float *A, *B, *C;
static int M, N, K;
static int num_threads;
static int mpi_rank, mpi_world_size; //taskid, numtasks

static int rows;
static int min(int x, int y) {
  return x < y ? x : y;
}

#define ITILESIZE (32)
#define JTILESIZE (1024)
#define KTILESIZE (1024)

static void mat_mul_omp() {
#pragma omp parallel num_threads(num_threads)
{

  int slice = rows / num_threads;
  int tid = omp_get_thread_num();
  int is = slice * tid       + min(tid    , rows % num_threads);
  int ie = slice * (tid + 1) + min(tid + 1, rows % num_threads);

  for (int ii = is; ii < ie; ii += ITILESIZE) {
    for (int jj = 0; jj < N; jj += JTILESIZE) {
      for (int kk = 0; kk < K; kk += KTILESIZE) {

        for (int k = kk; k < min(K, kk + KTILESIZE); k++) {
          for (int i = ii; i < min(ie, ii + ITILESIZE); i++) {
            float ar = A[i * K + k];
            for (int j = jj; j < min(N, jj + JTILESIZE); j+=1) {
              C[i * N + j] += ar * 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) {

  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
  int averow, extra, dest;
  int offset, src;
  int numworkers, mtype;
  MPI_Status status;
  MPI_Request request;

  numworkers = mpi_world_size-1;
  averow     = M/mpi_world_size;
  extra      = M%mpi_world_size;

  rows       = (mpi_rank == numworkers) ? averow+extra : averow;
  offset     = averow;


   // WORKER initialize
  if (mpi_rank > 0) {
//    printf("\nI_%d\n", mpi_rank);
    alloc_mat(&A, rows, K);
    alloc_mat(&B, K, N);
    alloc_mat(&C, rows, N);
    zero_mat(C, rows, N);
  }


    // MASTER send to WORKER 
if (mpi_rank == 0) {
  A = _A, B = _B, C = _C;  //order change upper->lower
    mtype  = FROM_MASTER;
  for (dest=1; dest<=numworkers; dest++)
  {
    rows   = (dest == numworkers) ? averow+extra           : averow;

    MPI_Isend(&offset     ,   1  ,   MPI_INT, dest, mtype, MPI_COMM_WORLD, &request);
    MPI_Isend(&rows       ,   1  ,   MPI_INT, dest, mtype, MPI_COMM_WORLD, &request);
    MPI_Isend(&A[offset*K],rows*K, MPI_FLOAT, dest, mtype, MPI_COMM_WORLD, &request);
    MPI_Isend(&B[0]       , K*N  , MPI_FLOAT, dest, mtype, MPI_COMM_WORLD, &request);
 
    offset += rows;
  }
}


  if (mpi_rank == 0) {
    mat_mul_omp();

  }

   // MASTER rcv from WORKER
if (mpi_rank == 0) {
  mtype = FROM_WORKER;
  for (src=1; src<=numworkers; src++)
  {
    MPI_Recv(&offset,   1,           MPI_INT, src , mtype, MPI_COMM_WORLD, &status);
    MPI_Recv(&rows    , 1,           MPI_INT, src , mtype, MPI_COMM_WORLD, &status);
    MPI_Recv(&C[offset*N], rows*N, MPI_FLOAT, src , mtype, MPI_COMM_WORLD, &status);
  }

  rows       = averow;
}






//MPI_Barrier(MPI_COMM_WORLD, &request);   /////////////////////////////////////////////////
    // WORKER rcv from MASTER 
if (mpi_rank > 0) {
    mtype = FROM_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[0]  , rows*K, MPI_FLOAT, MASTER , mtype, MPI_COMM_WORLD, &status);
    MPI_Recv(&B[0]  ,    K*N, MPI_FLOAT, MASTER , mtype, MPI_COMM_WORLD, &status);



#pragma omp parallel num_threads(num_threads)
{

  int slice = rows / num_threads;
  int tid = omp_get_thread_num();
  int is = slice * tid       + min(tid    , rows % num_threads);
  int ie = slice * (tid + 1) + min(tid + 1, rows % num_threads);

  for (int ii = is; ii < ie; ii += ITILESIZE) {
    for (int jj = 0; jj < N; jj += JTILESIZE) {
      for (int kk = 0; kk < K; kk += KTILESIZE) {

        for (int k = kk; k < min(K, kk + KTILESIZE); k++) {
          for (int i = ii; i < min(ie, ii + ITILESIZE); i++) {
            float ar = A[i * K + k];
            for (int j = jj; j < min(N, jj + JTILESIZE); j+=1) {
              C[i * N + j] += ar * B[k * N + j];
            }
          }
        }

      }
    }
  }
}







    // WORKER send to MASTER
    mtype = FROM_WORKER;
    MPI_Isend(&offset, 1     ,   MPI_INT, MASTER, mtype, MPI_COMM_WORLD, &request);
    MPI_Isend(&rows  , 1     ,   MPI_INT, MASTER, mtype, MPI_COMM_WORLD, &request);
    MPI_Isend(&C[0]  , rows*N, MPI_FLOAT, MASTER, mtype, MPI_COMM_WORLD, &request);
}








  // FIXME: for now, only root process runs the matrix multiplication.
  //if (mpi_rank == 0) 
  //  mat_mul_omp();

}