/*
 *  N: # of input sets
 *  C: # of channels
 *  H: input height
 *  W: input width
 *  K: # of filters -> # of output sets
 *  R: filter height
 *  S: filter width
 *
 *  OH: output height
 *  OW: output width
 */
#include "convolution.h"
#include "omp.h"
#include "util.h"
#include <mpi.h>
#include <stdio.h>
#include <fstream>

#define I_BSIZE 128
#define J_BSIZE 1024
#define K_BSIZE 1024

#define MIN(x, y) (((x) < (y)) ? (x) : (y))

static float *input, *output, *filter, *col_buf;
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 MPI_Status status;

void im2col_cpu(float *data_im, int batches, int channels, int height, int width,
            int kernel_h, int kernel_w, int pad, int stride,
            int dilation, float *data_col)
{
  register const long dil_kernel_h = (kernel_h - 1) * dilation + 1;
  register const long dil_kernel_w = (kernel_w - 1) * dilation + 1;
  register const long height_col = (height + 2 * pad - dil_kernel_h) / stride + 1;
  register const long width_col = (width + 2 * pad - dil_kernel_w) / stride + 1;
  register const long channels_col = channels * kernel_h * kernel_w;
  register const long batch_stride = channels_col * height_col * width_col;
  register const long sbatch_stride = C * H * W;

  #pragma omp parallel for num_threads(40)
  for (int b = 0; b < batches; ++b) {
    for (long c = 0; c < channels_col; ++c) {
      long w_offset = c % kernel_w;
      long h_offset = (c / kernel_w) % kernel_h;
      long c_im = c / kernel_h / kernel_w;

      const long hc0 = h_offset * dilation - pad;
      const long wc0 = w_offset * dilation - pad;
      for (long h = 0; h < height_col; ++h) {
        long h_pad = h * stride + hc0;

        const long row_offset = (c * height_col + h) * width_col;
        const long srow_offset = (c_im * height + h_pad) * width;
        for (long w = 0; w < width_col; ++w) {
          long w_pad = w * stride + wc0;
          if ((((unsigned long)h_pad) < ((unsigned long)height)) && (((unsigned long)w_pad) < ((unsigned long)width)))
            data_col[b * batch_stride + row_offset + w] = data_im[b * sbatch_stride + srow_offset + w_pad];
          else {
            data_col[b * batch_stride + row_offset + w] = 0.;
          }
        }
      }
    }
  }
}

/*
 * _A:  filter
 * _B:  input
 * _C:  output
 * _BA: batch size
 * _M:  K
 * _N:  OH * OW
 * _K:  C * R * S
 */
static void sgemm_cpu (float *_A, float *_B, float *_C, int _BA, int _M, int _N, int _K) 
{
  register const long batch_stride = _M * _N;
  register const long sbatch_stride = _K * _N;

#pragma omp parallel for num_threads(40) shared(_A, _B, _C)
  for (long b = 0; b < _BA; b ++)
  for (long ii = 0; ii < _M; ii += I_BSIZE)
    for (long jj = 0; jj < _N; jj += J_BSIZE)
      for (long kk = 0; kk < _K; kk += K_BSIZE)
        for (long k = kk; k < MIN(_K, kk + K_BSIZE); k++)
          for (long i = ii; i < MIN(_M, ii + I_BSIZE); i++) {
            float tmp = _A[i * _K + k];
            for (long j = jj; j < MIN(_N, jj + J_BSIZE); j++) {
              _C[b * batch_stride + i * _N + j] += tmp * _B[b * sbatch_stride + k * _N + j];
            }
          }
}

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) 
{
  input = _input;
  output = _output;
  filter = _filter;

  int size[mpi_world_size];
  int is[mpi_world_size], ie[mpi_world_size];

  for (int i=0; i<mpi_world_size; i++) {
      is[i] = N / mpi_world_size * i;
      ie[i] = N / mpi_world_size * (i+1);
  }

  ie[mpi_world_size-1] = N;

  for (int i=0; i<mpi_world_size; i++) {
      size[i] = ie[i] - is[i];
  }

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

  alloc_tensor(&col_buf, 1, size[mpi_rank], R * S * C, OH * OW);

  // Scatter
  if (mpi_rank == 0) {
    for (int i=1; i<mpi_world_size; i++) {
        MPI_Send(input + (is[i]*C*H*W), size[i]*C*H*W, MPI_FLOAT, i, 0, MPI_COMM_WORLD);
    }
  } else {
    alloc_tensor(&input, size[mpi_rank], C, H, W);
    alloc_tensor(&output, size[mpi_rank], K, OH, OW);
    alloc_tensor(&filter, K, C, R, S);
    MPI_Recv(input, size[mpi_rank]*C*H*W, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, &status);
  }

  // Broadcast
  MPI_Bcast(filter, K*C*R*S, MPI_FLOAT, 0, MPI_COMM_WORLD);

  // Calculation
  im2col_cpu(input, size[mpi_rank], C, H, W, R, S, pad, stride, dilation, col_buf);

  sgemm_cpu(filter, col_buf, output, size[mpi_rank], K, OH * OW, C * R * S);
//  convolution_omp(size[mpi_rank]);

  // Gather
  if (mpi_rank == 0) {
      for (int i=1; i<mpi_world_size; i++) {
          MPI_Recv(output + (is[i]*K*OH*OW), size[i]*K*OH*OW, MPI_FLOAT, i, 1, MPI_COMM_WORLD, &status);
      }
  } else {
      MPI_Send(output, size[mpi_rank]*K*OH*OW, MPI_FLOAT, 0, 1, MPI_COMM_WORLD);
  }
}

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);
}

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