#include "convolution.h"
#include <mpi.h>
#include <stdio.h>
#include <cuda_runtime.h>
#include "util.h"
#define MAX_NUM_GPU 4
#define TS 8

#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;
int num_devices = 4;

__global__ void conv(
    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 OH, OW;

  const int globalRow = blockDim.x * blockIdx.x + threadIdx.x;
  const int globalCol = blockDim.y * blockIdx.y + threadIdx.y;

  OH = (_H + 2 * _pad - _dilation * (_R - 1) - 1) / _stride + 1;
  OW = (_W + 2 * _pad - _dilation * (_S - 1) - 1) / _stride + 1;
  int n,k,w;
  w = globalCol;
  n = w/(_K*OH);
  w = w - n*(_K*OW);
  k = w / OW;
  w = w - k * OH;
  int col = w;
  int row = globalRow;

  if(globalRow >= OH || globalCol >= _N*_K*OW) return;
  int start_row = row * _stride - _pad;
  int start_col = col * _stride - _pad;

            float o = 0.f;
            for (int c = 0; c < _C; ++c) {
              for (int r = 0; r < _R; ++r) {
                for (int s = 0; s < _S; ++s) {
                  int h = start_row + r * _dilation;
                  int w = start_col + 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 + row * OW + col] = o;
	//    __syncthreads();
          
}
static float *in_d[MAX_NUM_GPU];
static float *out_d[MAX_NUM_GPU];
static float *filter_d[MAX_NUM_GPU];
static int Nbegin[MAX_NUM_GPU], Nend[MAX_NUM_GPU];

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

  int size[2];
  MPI_Request request;
  MPI_Status status;
  if(mpi_world_size == 2) size[1] = _N/2;
  else                    size[1] = 0;
  size[0] = _N-size[1];


  if (mpi_rank == 0 && mpi_world_size==2) {
          MPI_Isend(&_input[size[0] * _C * _H * _W], size[1] * _C * _H * _W, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &request);
          MPI_Isend(_filter, _K*_C*_R*_Rs, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &request);
  }
  else if(mpi_world_size==2){
          alloc_tensor(&_input, size[1], _C, _H, _W);
          alloc_tensor(&_output, size[1], _K, OH, OW);
          alloc_tensor(&_filter, _K, _C, _R, _Rs);
          MPI_Recv(_input, size[1]*_C*_H*_W, MPI_FLOAT, 0,0,MPI_COMM_WORLD, &status);
          MPI_Recv(_filter, _K*_C*_R*_Rs, MPI_FLOAT, 0,0, MPI_COMM_WORLD, &status);
  }
  // Upload A and B matrix to every GPU
  for (int i = 0; i < num_devices; i++) {
  	CUDA_CALL( cudaMemcpy(in_d[i], _input + (Nbegin[i]-size[1]*mpi_rank) * _C * _H * _W, (Nend[i]-Nbegin[i]) * _C * _H * _W * sizeof(float), cudaMemcpyHostToDevice) );
	CUDA_CALL( cudaMemcpy(filter_d[i], _filter, _K * _C * _R * _Rs * sizeof(float), cudaMemcpyHostToDevice) );
  }
  for(int i=0;i<num_devices;i++){
	dim3 gridDim((OH+TS-1)/TS, ((Nend[i]-Nbegin[i])*K*OW+TS-1)/TS,1);
	dim3 blockDim(TS, TS, 1);

	CUDA_CALL(cudaSetDevice(i));
	conv<<<gridDim, blockDim>>>(in_d[i], out_d[i], filter_d[i], (Nend[i]-Nbegin[i]), _C, _H, _W, _K, _R, _Rs, _pad, _dilation, _stride);				
  }

  for (int i = 0; i < num_devices; i++) {
  	CUDA_CALL ( cudaMemcpy( _output+(Nbegin[i]-size[1]*mpi_rank)* _K * OH * OW, out_d[i], (Nend[i]-Nbegin[i]) * _K * OH * OW*sizeof(float), cudaMemcpyDeviceToHost));

  }
  if(mpi_rank == 0 && mpi_world_size==2){
	MPI_Recv(&_output[size[0]*_K*OH*OW], size[1]*_K*OH*OW, MPI_FLOAT, 1, 0, MPI_COMM_WORLD, &status);
  }
  else if(mpi_world_size==2){
	MPI_Isend(_output, size[1]*_K*OH*OW, MPI_FLOAT, 0, 0, MPI_COMM_WORLD,&request);
  }

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;

  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;

  CUDA_CALL( cudaGetDeviceCount(&num_devices) );

  printf("Using %d devices\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);
  }


  if(mpi_rank==0){
  	for (int i = 0; i < num_devices; i++) {
  		Nbegin[i] = (N / num_devices) * i/mpi_world_size;
  		Nend[i] = (N / num_devices) * (i + 1)/mpi_world_size;
  	}
  	Nend[num_devices - 1] = N/mpi_world_size;
  }
  else {
	for(int i=0;i<num_devices;i++){
		Nbegin[i] = (N / num_devices) * i/mpi_world_size + N/mpi_world_size;
		Nend[i] = (N / num_devices) * (i+1)/mpi_world_size + N/mpi_world_size;
	}
  	Nend[num_devices - 1] = N;
  }
  // Allocate device memory for each GPU
  for (int i = 0; i < num_devices; i++) {
    CUDA_CALL( cudaSetDevice(i) );
    CUDA_CALL( cudaMalloc(&in_d[i], (Nend[i]-Nbegin[i]) * C * H * W * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&out_d[i], (Nend[i]-Nbegin[i]) * K * OH * OW * sizeof(float)) );
    CUDA_CALL( cudaMalloc(&filter_d[i], K * C * R * S * sizeof(float)) );
  }
  

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

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

}