#include <iostream>
#include <fstream>
#include <cmath>
#include <chrono>

#define dim_x 128
#define dim_fc1 32
#define batch_sz 8

using namespace std;

#include <cublas_v2.h>

#define CHECK(value) {                                          \
    cudaError_t _m_cudaStat = value;                                        \
    if (_m_cudaStat != cudaSuccess) {                                       \
        cout<< "Error:" << cudaGetErrorString(_m_cudaStat) \
            << " at line " << __LINE__ << " in file " << __FILE__ << "\n"; \
        exit(1);                                                            \
    } }

#define cublasErrCheck(stat) { cublasErrCheck_((stat), __FILE__, __LINE__); }
void cublasErrCheck_(cublasStatus_t stat, const char *file, int line) {
   if (stat != CUBLAS_STATUS_SUCCESS) {
      fprintf(stderr, "cuBLAS Error: %d %s %d\n", stat, file, line);
   }
}

#include "sgemm.c"

const float one=1, zero=0, two=2;

float frand()  { return ((double) rand()) / RAND_MAX * 2 - 1; }

void init_random_w(float *w1, float *w2)
{
  for (int i = 0; i < dim_x; i++)
    for (int j = 0; j < dim_fc1; j++)
      w1[i + j*dim_x] = frand();
  for (int j = 0; j < dim_fc1; j++)
    w2[j] = frand();
}


// модель для CPU
class CpuModel
{
public:
  float *w1, *w2, *x, *sigmoid, *fc2, *diff, *dsigma_mul_w2_mul_x;

  CpuModel(bool random)
  {
    w1 = new float[dim_x*dim_fc1];   // shape = 128 x 32
    w2 = new float[dim_fc1];
    x = NULL;
    fc2 = NULL;
    // выделяем заранее память под результаты промежуточных вычислений
    sigmoid = new float[batch_sz*dim_fc1]; // shape = 8 x 32
    diff = new float[batch_sz];
    dsigma_mul_w2_mul_x = new float[batch_sz*dim_x*dim_fc1];  // shape = 8 x 128 x 32

    if (random)  init_random_w(w1, w2);
    else
    {
      // считываем веса из файла
      ifstream f1("w1.txt");
      for (int i = 0; i < dim_x; i++)
        for (int j = 0; j < dim_fc1; j++)
          // располагаем в памяти массивы как в фортране - по столбцам
          // тогда элементы сворачиваемой размерности будут расположены вплотную друг ко другу
          // и функция перемножения матриц сможет использовать тензорные ядра
          f1 >> w1[i + j*dim_x];
      ifstream f2("w2.txt");
      for (int j = 0; j < dim_fc1; j++)
        f2 >> w2[j];
    }
  }

  void forward(float *x_, float *y)
  {
    x = x_;
    // fc1 = np.dot(x, w1)
    sgemm("N", "N", batch_sz, dim_fc1, dim_x, &one, x, batch_sz, w1, dim_x, &zero, sigmoid, batch_sz);
    for (int i = 0; i < batch_sz*dim_fc1; i++) sigmoid[i] = 1.0f/(1.0f+expf(-sigmoid[i]));
    fc2 = y;
    // fc2 = np.dot(sigmoid, w2)
    sgemm("N", "N", batch_sz, 1, dim_fc1, &one, sigmoid, batch_sz, w2, dim_fc1, &zero, fc2, batch_sz);
  }

  void backward(float *y_true, float *grad_loss_by_w1, float *grad_loss_by_w2)
  {
    for (int i=0; i<batch_sz; i++)  diff[i] = fc2[i] - y_true[i]; // shape = 8 x 1

    // вариант программы с функцией gemm, использующей (на видеокарте) тензорные ядра
    /*
    for (int i=0; i<batch_sz; i++)
      for (int j=0; j<dim_x; j++)
        for (int k=0; k<dim_fc1; k++)
        {
          float s = sigmoid[i + batch_sz*k];
          float dsigma = s*(1-s);
          float dsigma_mul_w2 = dsigma * w2[k];
          dsigma_mul_w2_mul_x[i + batch_sz*j + batch_sz*dim_x*k] = dsigma_mul_w2 * x[i + batch_sz*j];
        }
    // grad_loss_by_w1 = 2*np.sum(dsigma_mul_w2_mul_x * diff[:,:,np.newaxis], axis=0)  # shape = 128 x 32
    sgemm("N", "N", 1, dim_x*dim_fc1, batch_sz, &two, diff, 1, dsigma_mul_w2_mul_x, batch_sz, &zero, grad_loss_by_w1, 1);
    */

    // другой вариант без умножения матриц на тензорных ядрах, но с экономией памяти - не используется
    // трехмерная матрица dsigma_mul_w2_mul_x

    for (int j=0; j<dim_x; j++)
      for (int k=0; k<dim_fc1; k++)
      {
        float sum = 0;
        for (int i=0; i<batch_sz; i++)
        {
          float s = sigmoid[i + batch_sz*k];
          float dsigma = s*(1-s);
          float dsigma_mul_w2 = dsigma * w2[k];
          sum += dsigma_mul_w2 * x[i + batch_sz*j] * diff[i];
        }
        grad_loss_by_w1[j + dim_x*k] = 2*sum;
      }

    // grad_loss_by_w2 = 2*np.dot(self.sigmoid.T, diff)
    sgemm("N", "N", 1, dim_fc1, batch_sz, &two, diff, 1, sigmoid, batch_sz, &zero, grad_loss_by_w2, 1);
  }

  ~CpuModel()
  {
    delete[] w1;
    delete[] w2;
    delete[] sigmoid;
    delete[] diff;
    delete[] dsigma_mul_w2_mul_x;
  }
};



__global__ void sigmoid_kernel(float* tau, int sz)
{
  int i = blockDim.x * blockIdx.x + threadIdx.x;
  if (i < sz) tau[i] = 1.0f / (1.0f + exp (-tau[i]));
}

__global__ void diff_kernel(float* a, float *b, float *c, int sz)
{
  int i = blockDim.x * blockIdx.x + threadIdx.x;
  if (i < sz) c[i] = a[i] - b[i];
}

__global__ void variant1_kernel(float* sigmoid, float *x, float *w2, float *diff, float *dsigma_mul_w2_mul_x)
{
  int j = blockDim.x * blockIdx.x + threadIdx.x;
  int i = blockDim.y * blockIdx.y + threadIdx.y;
  int k = blockDim.z * blockIdx.z + threadIdx.z;
  if (i < batch_sz && j < dim_x && k < dim_fc1)
  {
    float s = sigmoid[i + batch_sz*k];
    float dsigma = s*(1-s);
    float dsigma_mul_w2 = dsigma * w2[k];
    dsigma_mul_w2_mul_x[i + batch_sz*j + batch_sz*dim_x*k] = dsigma_mul_w2 * x[i + batch_sz*j];
  }
}

__global__ void variant2_kernel(float* sigmoid, float *x, float *w2, float *diff, float *grad_loss_by_w1)
{
  int j = blockDim.x * blockIdx.x + threadIdx.x;
  int k = blockDim.y * blockIdx.y + threadIdx.y;
  if (j < dim_x && k < dim_fc1)
  {
    float sum = 0;
    for (int i=0; i<batch_sz; i++)
    {
      // TODO: подумайте, как можно было бы ускорить эту функцию, в чем здесь проблемы
      float s = sigmoid[i + batch_sz*k];
      float dsigma = s*(1-s);
      float dsigma_mul_w2 = dsigma * w2[k];
      sum += dsigma_mul_w2 * x[i + batch_sz*j] * diff[i];
    }
    grad_loss_by_w1[j + dim_x*k] = 2*sum;
  }
}

// модель для GPU
class GpuModel
{
public:
  // все это указатели на видеопамять!!!!
  float *w1, *w2, *x, *sigmoid, *fc2, *diff, *dsigma_mul_w2_mul_x;

  cublasHandle_t cublasHandle;

  GpuModel()
  {
    CHECK( cudaMalloc(&w1, dim_x*dim_fc1*sizeof(float)) );
    CHECK( cudaMalloc(&w2, dim_fc1*sizeof(float)) );
    x = NULL;
    fc2 = NULL;
    // выделяем заранее память под результаты промежуточных вычислений
    CHECK( cudaMalloc(&sigmoid, batch_sz*dim_fc1*sizeof(float)) );
    CHECK( cudaMalloc(&diff, batch_sz*sizeof(float)) );
    CHECK( cudaMalloc(&dsigma_mul_w2_mul_x, batch_sz*dim_x*dim_fc1*sizeof(float)) );
    float *w1_host = new float[dim_x*dim_fc1];
    float *w2_host = new float[dim_fc1];
    init_random_w(w1_host, w2_host);
    CHECK( cudaMemcpy(w1, w1_host, dim_x*dim_fc1*sizeof(float),cudaMemcpyHostToDevice) );
    CHECK( cudaMemcpy(w2, w2_host, dim_fc1*sizeof(float),cudaMemcpyHostToDevice) );
    delete[] w1_host;
    delete[] w2_host;
    cublasErrCheck(cublasCreate(&cublasHandle));
    cublasErrCheck(cublasSetMathMode(cublasHandle, CUBLAS_TENSOR_OP_MATH));
  }

  void forward(float *x_, float *y)
  {
    x = x_;
    // fc1 = np.dot(x, w1)
    cublasErrCheck( cublasGemmEx(cublasHandle, CUBLAS_OP_N, CUBLAS_OP_N, batch_sz, dim_fc1, dim_x,
      &one, x, CUDA_R_32F, batch_sz, w1, CUDA_R_32F, dim_x,
      &zero, sigmoid, CUDA_R_32F, batch_sz, CUDA_R_32F, CUBLAS_GEMM_DFALT_TENSOR_OP) ); // все требования для параметров выполнены, поэтому CUBLAS_GEMM_DFALT_TENSOR_OP
    sigmoid_kernel<<<(batch_sz*dim_fc1+255)/256, 256>>>(sigmoid, batch_sz*dim_fc1);
    fc2 = y;
    // fc2 = np.dot(sigmoid, w2)
    cublasErrCheck( cublasGemmEx(cublasHandle, CUBLAS_OP_N, CUBLAS_OP_N, batch_sz, 1, dim_fc1,
      &one, sigmoid, CUDA_R_32F, batch_sz, w2, CUDA_R_32F, dim_fc1,
      &zero, fc2, CUDA_R_32F, batch_sz, CUDA_R_32F, CUBLAS_GEMM_DFALT_TENSOR_OP) ); // все требования для параметров выполнены, поэтому CUBLAS_GEMM_DFALT_TENSOR_OP
  }

  void backward(float *y_true, float *grad_loss_by_w1, float *grad_loss_by_w2)
  {
    diff_kernel<<<(batch_sz+255)/256, 256>>>(fc2, y_true, diff, batch_sz);

    // вариант программы с функцией gemm, использующей (на видеокарте) тензорные ядра
    /*
    variant1_kernel<<<dim3((dim_x+7)/8, (batch_sz+31)/32, (dim_fc1+3)/4), dim3(8,32,4)>>>(sigmoid, x, w2, diff, dsigma_mul_w2_mul_x);
    // grad_loss_by_w1 = 2*np.sum(dsigma_mul_w2_mul_x * diff[:,:,np.newaxis], axis=0)
    cublasErrCheck( cublasGemmEx(cublasHandle, CUBLAS_OP_N, CUBLAS_OP_N, 1, dim_x*dim_fc1, batch_sz,
      &two, diff, CUDA_R_32F, 1, dsigma_mul_w2_mul_x, CUDA_R_32F, batch_sz,
      &zero, grad_loss_by_w1, CUDA_R_32F, 1, CUDA_R_32F, CUBLAS_GEMM_DEFAULT) ); // первая размерность не кратна 4, поэтому CUBLAS_GEMM_DEFAULT
    */

    // другой вариант без умножения матриц на тензорных ядрах, но с экономией памяти - не используется
    // трехмерная матрица dsigma_mul_w2_mul_x
    variant2_kernel<<<dim3((dim_x+15)/16, (dim_fc1+15)/16, 1), dim3(16,16,1)>>>(sigmoid, x, w2, diff, grad_loss_by_w1);

    // grad_loss_by_w2 = 2*np.dot(self.sigmoid.T, diff)
    cublasErrCheck( cublasGemmEx(cublasHandle, CUBLAS_OP_N, CUBLAS_OP_N, 1, dim_fc1, batch_sz,
      &two, diff, CUDA_R_32F, 1, sigmoid, CUDA_R_32F, batch_sz,
      &zero, grad_loss_by_w2, CUDA_R_32F, 1, CUDA_R_32F, CUBLAS_GEMM_DEFAULT) ); // первая размерность не кратна 4, поэтому CUBLAS_GEMM_DEFAULT
  }

  ~GpuModel()
  {
    cudaFree(w1);
    cudaFree(w2);
    cudaFree(sigmoid);
    cudaFree(diff);
    cudaFree(dsigma_mul_w2_mul_x);
  }
};


// Сравнение работы нашей CPU модели и python-модели. Работает только если размеры в python
// те же, что и здесь.
// Обычно если CPU модель правильно работает для одного размера, то
// правильно и для других размеров. Для GPU это неверно
void check_with_python()
{
  CpuModel check_model(false);

  float *x = new float[10000*dim_x];
  ifstream fx("x.txt");
  for (int i=0; i < 10000; i++)
    for (int j=0; j < dim_x; j++)
      fx >> x[i + 10000*j];
  float *y0 = new float[10000];
  ifstream fy("y0.txt");
  for (int i=0; i < 10000; i++)  fy >> y0[i];

  int *ind = new int[batch_sz];
  ifstream find("ind.txt");
  for (int i=0; i < batch_sz; i++)  find >> ind[i];

  float *x_batch = new float[batch_sz*dim_x];
  float *y_true = new float[batch_sz];
  for (int i=0; i < batch_sz; i++)
  {
    for (int j=0; j < dim_x; j++)
      x_batch[i + batch_sz*j] = x[ind[i] + 10000*j];
    y_true[i] = y0[ind[i]];
  }

  // проверка прямого хода
  float *y_pred = new float[batch_sz];
  check_model.forward(x_batch, y_pred);
  cout << "y_pred = ";
  for (int i=0; i < batch_sz; i++) cout << y_pred[i] << "\n";

  // проверка обратного хода
  float *grad_loss_by_w1 = new float[dim_x*dim_fc1];
  float *grad_loss_by_w2 = new float[dim_fc1];
  check_model.backward(y_true, grad_loss_by_w1, grad_loss_by_w2);
  cout << "grad_loss_by_w1 = ";
  for (int i=0; i<3; i++)
  {
    for (int j=0; j<3; j++)
      cout << grad_loss_by_w1[i + dim_x*j] << " ";
    cout << " ... ";
    for (int j=dim_fc1-3; j<dim_fc1; j++)
      cout << grad_loss_by_w1[i + dim_x*j] << " ";
    cout << "\n";
  }
  cout << "........................................\n";
  for (int i=dim_x-3; i<dim_x; i++)
  {
    for (int j=0; j<3; j++)
      cout << grad_loss_by_w1[i + dim_x*j] << " ";
    cout << " ... ";
    for (int j=dim_fc1-3; j<dim_fc1; j++)
      cout << grad_loss_by_w1[i + dim_x*j] << " ";
    cout << "\n";
  }
  cout << "grad_loss_by_w2 = ";
  for (int i=0; i<dim_fc1; i++) cout << grad_loss_by_w2[i] << "\n";
}


// сравнивает результаты cpu и gpu версий нейронной сети
// работает для любых раметров слоев
void check_by_compare_cpu_gpu()
{
  // для замеров времени
  cudaEvent_t startCUDA, stopCUDA;
  cudaEventCreate(&startCUDA);
  cudaEventCreate(&stopCUDA); 
  float elapsedTimeCUDA;
  
  srand(1);
  CpuModel cpu_model(true);
  srand(1);
  GpuModel gpu_model;

  float *x_batch = new float[batch_sz*dim_x];
  for (int i=0; i<batch_sz*dim_x; i++)  x_batch[i] = frand();
  float *y_pred_cpu = new float[batch_sz];
  
  auto startCPU = chrono::system_clock::now();
  cpu_model.forward(x_batch, y_pred_cpu);
  double elapsedTimeCPU = chrono::duration_cast<std::chrono::microseconds>(chrono::system_clock::now() - startCPU).count() / 1000.0;

  float *x_batch_dev, *y_pred_gpu_dev;
  CHECK( cudaMalloc(&x_batch_dev, batch_sz*dim_x*sizeof(float)) );
  CHECK( cudaMemcpy(x_batch_dev, x_batch, batch_sz*dim_x*sizeof(float), cudaMemcpyHostToDevice) );
  CHECK( cudaMalloc(&y_pred_gpu_dev, batch_sz*sizeof(float)) );
  
  cudaEventRecord(startCUDA,0);
  gpu_model.forward(x_batch_dev, y_pred_gpu_dev);
  cudaEventRecord(stopCUDA,0);
  cudaEventSynchronize(stopCUDA);
  CHECK(cudaGetLastError());
  cudaEventElapsedTime(&elapsedTimeCUDA, startCUDA, stopCUDA);

  // check
  float *y_pred_gpu_host = new float[batch_sz];
  CHECK( cudaMemcpy(y_pred_gpu_host, y_pred_gpu_dev, batch_sz*sizeof(float), cudaMemcpyDeviceToHost) );
  for (int i = 0; i < batch_sz; i++)
    if (abs(y_pred_cpu[i] - y_pred_gpu_host[i]) > 1e-3)
    {
      cout << "Error in forward. Element N " << i << ": y_pred_cpu[i] = " << y_pred_cpu[i]
           << " y_pred_gpu_host[i] = " << y_pred_gpu_host[i] << "\n";
      exit(1);
    }
  cout << "GPU forward - OK\n";
  cout << "CPU time = " << elapsedTimeCPU << " ms\n";
  cout << "CUDA time = " << elapsedTimeCUDA << " ms\n";

  // проверка обратного хода
  float *y_true = new float[batch_sz];
  for (int i=0; i<batch_sz; i++)  y_true[i] = frand();
  float *grad_loss_by_w1_cpu = new float[dim_x*dim_fc1];
  float *grad_loss_by_w2_cpu = new float[dim_fc1];
  
  startCPU = chrono::system_clock::now();
  cpu_model.backward(y_true, grad_loss_by_w1_cpu, grad_loss_by_w2_cpu);
  elapsedTimeCPU = chrono::duration_cast<std::chrono::microseconds>(chrono::system_clock::now() - startCPU).count() / 1000.0;

  float *y_true_dev, *grad_loss_by_w1_gpu_dev, *grad_loss_by_w2_gpu_dev;
  CHECK( cudaMalloc(&y_true_dev, batch_sz*sizeof(float)) );
  CHECK( cudaMemcpy(y_true_dev, y_true, batch_sz*sizeof(float), cudaMemcpyHostToDevice) );
  CHECK( cudaMalloc(&grad_loss_by_w1_gpu_dev, dim_x*dim_fc1*sizeof(float)) );
  CHECK( cudaMalloc(&grad_loss_by_w2_gpu_dev, dim_fc1*sizeof(float)) );
  
  cudaEventRecord(startCUDA,0);
  gpu_model.backward(y_true_dev, grad_loss_by_w1_gpu_dev, grad_loss_by_w2_gpu_dev);
  cudaEventRecord(stopCUDA,0);
  cudaEventSynchronize(stopCUDA);
  CHECK(cudaGetLastError());
  cudaEventElapsedTime(&elapsedTimeCUDA, startCUDA, stopCUDA);

  // в w2 обычно меньше ошибок, чем в w1
  float *grad_loss_by_w2_gpu_host = new float[dim_fc1];
  CHECK( cudaMemcpy(grad_loss_by_w2_gpu_host, grad_loss_by_w2_gpu_dev, dim_fc1*sizeof(float), cudaMemcpyDeviceToHost) );
  for (int i = 0; i < dim_fc1; i++)
    if (abs(grad_loss_by_w2_cpu[i] - grad_loss_by_w2_gpu_host[i]) > 1e-3)
    {
      cout << "Error in backward grad_w2. Element N " << i << ": grad_loss_by_w2_cpu[i] = " << grad_loss_by_w2_cpu[i]
           << " grad_loss_by_w2_gpu_host[i] = " << grad_loss_by_w2_gpu_host[i] << "\n";
      exit(1);
    }
  cout << "GPU backward grad_w2 - OK\n";
  float *grad_loss_by_w1_gpu_host = new float[dim_x*dim_fc1];
  CHECK( cudaMemcpy(grad_loss_by_w1_gpu_host, grad_loss_by_w1_gpu_dev, dim_x*dim_fc1*sizeof(float), cudaMemcpyDeviceToHost) );
  for (int i = 0; i < dim_x; i++)
    for (int j = 0; j < dim_fc1; j++)
      if (abs(grad_loss_by_w1_cpu[i+dim_x*j] - grad_loss_by_w1_gpu_host[i+dim_x*j]) > 1e-3)
      {
          cout << "Error in backward grad_w1. Element N " << i << " " << j
            << ": grad_loss_by_w1_cpu[i,j] = " << grad_loss_by_w1_cpu[i+dim_x*j]
            << " grad_loss_by_w1_gpu_host[i,j] = " << grad_loss_by_w1_gpu_host[i+dim_x*j] << "\n";
          exit(1);
      }
  cout << "GPU backward grad_w1 - OK\n";
  cout << "CPU time = " << elapsedTimeCPU << " ms\n";
  cout << "CUDA time = " << elapsedTimeCUDA << " ms\n";
}

int main()
{
  check_with_python();
  // эту функцию нужно вызывать несколько раз, пробуя маленькие и большие размеры dim_x, batch_sz, dim_fc1
  check_by_compare_cpu_gpu();
}