#include "convolution.h"
#include <mpi.h>
#include <stdio.h>
#include <cuda_runtime.h>
#include "util.h"

#define MAX_NODE (4)
#define MAX_THREADS (80)
#define MAX_NUM_GPU 4
#define TS 8

#define MATRIX_SEND_DATA_MSG_ID 1000
#define MATRIX_SEND_RESULT_MSG_ID 1001

#define CUDA_CALL(f)                                                           \
  {                                                                            \
    cudaError_t err = (f);                                                     \
    if (err != cudaSuccess) {                                                  \
      fprintf(stderr, "CUDA error at [%s:%d] %d %s\n", __FILE__, __LINE__,     \
              err, cudaGetErrorString(err));                                   \
      exit(1);                                                                 \
    }                                                                          \
  }

static float *input, *output, *filter;
static int N, C, H, W;
static int K, R, S;
static int OH, OW;
static int pad;
static int dilation;
static int stride;
static int mpi_rank, mpi_world_size;

static int num_devices = 0;
static int startN[MAX_NODE], endN[MAX_NODE];
static int Nbegin[MAX_NUM_GPU], Nend[MAX_NUM_GPU];

static float *input_d[MAX_NUM_GPU];
static float *filter_d[MAX_NUM_GPU];
static float *output_d[MAX_NUM_GPU];

__global__ void run_convolution(
    float *_input, float *_output, float* _filter,
    int _N, int _C, int _H, int _W, int _K, int _R, int _S,
	int _stride, int _pad, int _dilation) {

    int _OH = (_H + 2 * _pad - _dilation * (_R - 1) - 1) / _stride + 1;
    int _OW = (_W + 2 * _pad - _dilation * (_S - 1) - 1) / _stride + 1;

    int n = blockIdx.x;
    int k = blockIdx.y;
    int oh = blockIdx.z;
    int ow = threadIdx.x;

    float o = 0.0f;
    for (int c = 0; c < _C; c++) {
      for (int r = 0; r < _R; r++) {
        for (int s = 0; s < _S; s++) {
          int h = oh * _stride - _pad + r * _dilation;
          int w = ow * _stride - _pad + s * _dilation;
          if (h < 0 || h >= _H || w < 0 || w >= _W) continue;
          float i = _input[n*_C*_H*_W + c*_H*_W + h*_W + w];
          float f = _filter[k*_C*_R*_S + c*_R*_S + r*_S + s];
          o += i * f;
        }
      }
    }
    _output[n*_K*_OH*_OW + k*_OH*_OW + oh*_OW + ow] = o;
}

void convolution(
    float *_input, float *_output, float *_filter, 
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {

  int input_size, filter_size, output_size;
  int i, m_size, sizeN = 0;
  MPI_Status status;
  MPI_Request req1[MAX_NODE], req2[MAX_NODE], req3[MAX_NODE], req4[MAX_NODE];

  input_size = C * H * W;
  filter_size = K * C * R * S;
  output_size = K * OH * OW;
  
  if (mpi_world_size <= mpi_rank)
    return;
  
  if (mpi_rank == 0) {
    input = _input;
    output = _output;
    filter = _filter;

    // Send Matrix Information
    for (i = 1; i < mpi_world_size; i++) {
      m_size = (endN[i] - startN[i]);
      MPI_Isend(&m_size, 1, MPI_INT, i, MATRIX_SEND_DATA_MSG_ID, MPI_COMM_WORLD, &req1[i]);
      MPI_Isend(&input[startN[i] * input_size], m_size * input_size, MPI_FLOAT, i, MATRIX_SEND_DATA_MSG_ID, MPI_COMM_WORLD, &req2[i]);
      MPI_Isend(&filter[0], filter_size, MPI_FLOAT, i, MATRIX_SEND_DATA_MSG_ID, MPI_COMM_WORLD, &req3[i]);
    }
    sizeN = endN[0];

    for (i = 0; i < num_devices; i++) {
      CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * input_size,
                 (Nend[i] - Nbegin[i]) * input_size * sizeof(float),
                 cudaMemcpyHostToDevice) );
      CUDA_CALL( cudaMemcpy(filter_d[i], filter, filter_size * sizeof(float), cudaMemcpyHostToDevice) );
    }

    // wait for sending
    for (i = 1; i < mpi_world_size; i++) {
      MPI_Wait(&req1[i], &status);
      MPI_Wait(&req2[i], &status);
      MPI_Wait(&req3[i], &status);
    }

  } else {
    MPI_Recv(&sizeN, 1, MPI_INT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
    m_size = sizeN;
    MPI_Irecv(&input[0], m_size * input_size, MPI_FLOAT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &req2[0]);
    MPI_Irecv(&filter[0], filter_size, MPI_FLOAT, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &req3[0]);
    // wait for receiving
    MPI_Wait(&req2[0], &status);
    MPI_Wait(&req3[0], &status);

    for (i = 0; i < num_devices; i++) {
      CUDA_CALL( cudaMemcpy(input_d[i], input + Nbegin[i] * input_size,
                 (Nend[i] - Nbegin[i]) * input_size * sizeof(float),
                 cudaMemcpyHostToDevice) );
      CUDA_CALL( cudaMemcpy(filter_d[i], filter, filter_size * sizeof(float), cudaMemcpyHostToDevice) );
    }
  }

  for (int i = 0; i < num_devices; i++) {
    m_size = (Nend[i] - Nbegin[i]);
	CUDA_CALL( cudaSetDevice(i) );
    dim3 blockDim(OW, 1, 1);
    dim3 gridDim(m_size, K, OH);

    //run_convolution<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], m_size, C, H, W, K, R, S, OH, OW, stride, pad, dilation);
    run_convolution<<<gridDim, blockDim>>>(input_d[i], output_d[i], filter_d[i], m_size, C, H, W, K, R, S, stride, pad, dilation);
  }

  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  }

  // Download C matrix from GPUs
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaMemcpy(output + Nbegin[i] * output_size, output_d[i],
               (Nend[i] - Nbegin[i]) * output_size * sizeof(float),
               cudaMemcpyDeviceToHost) );
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  }

  if (mpi_rank == 0) {

    // receiving the result
    for (i = 1; i < mpi_world_size; i++) {
      
      m_size = (endN[i] - startN[i]) * output_size;
      MPI_Irecv(&output[startN[i] * output_size], m_size, MPI_FLOAT, i, MPI_ANY_TAG, MPI_COMM_WORLD, &req4[i]);
    }

    // wait for receiving
    for (i = 1; i < mpi_world_size; i++) {
      MPI_Wait(&req4[i], &status);
    }
  } else {
    // sending the result
    m_size = (endN[mpi_rank] - startN[mpi_rank]) * output_size;
    MPI_Isend(&output[0], m_size, MPI_FLOAT, 0, MATRIX_SEND_RESULT_MSG_ID, MPI_COMM_WORLD, &req4[0]);
    MPI_Wait(&req4[0], &status);
  }
}

void cuda_init() {

  CUDA_CALL( cudaGetDeviceCount(&num_devices) );

  if (mpi_world_size <= 1) {
    if (N < 4) {
      num_devices = N;
    }
  }
  else {
    if (num_devices*2 > N) {
      num_devices /= 2;
    }
  }

  //printf("num device[%d]\n", num_devices);
  for (int i = 0; i < num_devices; i++) {
    cudaDeviceProp prop;
    CUDA_CALL( cudaGetDeviceProperties(&prop, i) );

    // Try printing more detailed information here
    printf("[GPU %d] %s\n", i, prop.name);
  }

  if (num_devices <= 0) {
    printf("No CUDA device found. Aborting\n");
    exit(1);
  }

  int m_size;
  int slice = (endN[mpi_rank] - startN[mpi_rank]) / num_devices;
  // Setup problem size for each GPU
  for (int i = 0; i < num_devices; i++) {
    Nbegin[i] = slice * i;
    Nend[i] = (i == num_devices - 1) ? (endN[mpi_rank] - startN[mpi_rank]) : slice * (i + 1);
  }

  // Allocate device memory for each GPU
  for (int i = 0; i < num_devices; i++) {
    m_size = Nend[i] - Nbegin[i];
    CUDA_CALL( cudaSetDevice(i) );
    // 
    CUDA_CALL( cudaMalloc(&input_d[i], m_size * C * H * W * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&filter_d[i], K * C * R * S * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&output_d[i], m_size * K * OH * OW * sizeof(float)) );
  }

  // DO NOT REMOVE; NEEDED FOR TIME MEASURE
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaDeviceSynchronize() );
  }
}

void convolution_init(
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
  N = _N; C = _C; H = _H; W = _W;
  K = _K; R = _R; S = _S;
  pad = _pad;
  dilation = _dilation;
  stride = _stride;
  int i, slice = 0, m_size;

  MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
  MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size);

  OH = (H + 2 * pad - dilation * (R - 1) - 1) / stride + 1;
  OW = (W + 2 * pad - dilation * (S - 1) - 1) / stride + 1;

  if (N < 8) {
    mpi_world_size = 1;
  }

  slice = N / mpi_world_size;
  
  for (i = 0; i < mpi_world_size; i++) {
    startN[i] = i * slice;
    endN[i] = (i == mpi_world_size - 1) ? N : (i + 1) * slice;
	m_size = endN[i] - startN[i];
    if (i != 0) {
      alloc_tensor(&input, m_size, C, H, W);
	  alloc_tensor(&output, m_size, K, OH, OW);
	  alloc_tensor(&filter, K, C, R, S);
    }
  }

  if (mpi_world_size <= mpi_rank)
    return;

  cuda_init();
}

void convolution_final(
    int _N, int _C, int _H, int _W,
    int _K, int _R, int _S,
    int _pad, int _dilation, int _stride) {
}