//////////////////////////////////////////////////////////////////////
// HIP version of Monte Carlo algorithm using hipRAND device API
//////////////////////////////////////////////////////////////////////

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>

#include <hip/hip_runtime.h>
#include <hiprand/hiprand.h>
#include <hiprand/hiprand_kernel.h>

#include "helper_hip.h"

//////////////////////////////////////////////////////////////////////
// HIP global constants
//////////////////////////////////////////////////////////////////////

__constant__ int N;
__constant__ float T, r, sigma, rho, alpha, dt, con1, con2;

//////////////////////////////////////////////////////////////////////
// kernel routines
//////////////////////////////////////////////////////////////////////

__global__ void setup_kernel(hiprandState_t *state) {
  const int id = threadIdx.x + blockIdx.x * blockDim.x;
  hiprand_init(1234ULL, id, 0ULL, &state[id]);
}

__global__ void pathcalc(float *d_v, hiprandState_t *state, int npaths,
                         int totalThreads) {
  const int id = threadIdx.x + blockIdx.x * blockDim.x;
  hiprandState_t localState = state[id];

  for (int path = id; path < npaths; path += totalThreads) {
    float s1 = 1.0f;
    float s2 = 1.0f;

    for (int n = 0; n < N; n++) {
      const float y1 = hiprand_normal(&localState);
      const float y2 = rho * y1 + alpha * hiprand_normal(&localState);
      s1 = s1 * (con1 + con2 * y1);
      s2 = s2 * (con1 + con2 * y2);
    }

    float payoff = 0.0f;
    if (fabsf(s1 - 1.0f) < 0.1f && fabsf(s2 - 1.0f) < 0.1f) {
      payoff = expf(-r * T);
    }
    d_v[path] = payoff;
  }

  state[id] = localState;
}

//////////////////////////////////////////////////////////////////////
// Main program
//////////////////////////////////////////////////////////////////////

int main(int argc, const char **argv) {
  const int NPATH = 9600000;
  int h_N = 100;
  float h_T, h_r, h_sigma, h_rho, h_alpha, h_dt, h_con1, h_con2;
  float *h_v = nullptr;
  float *d_v = nullptr;
  hiprandState_t *state = nullptr;
  double sum1, sum2;

  // initialise card
  findHipDevice(argc, argv);

  // initialise HIP timing
  float milli = 0.0f;
  hipEvent_t start, stop;
  checkHipErrors(hipEventCreate(&start));
  checkHipErrors(hipEventCreate(&stop));

  // choose block size and use occupancy to determine how many blocks can
  // execute simultaneously without queueing
  int blockSize = 256;
  int maxBlocksPerSM = 0;
  int device = 0;
  hipDeviceProp_t prop{};
  checkHipErrors(hipGetDevice(&device));
  checkHipErrors(hipGetDeviceProperties(&prop, device));
  checkHipErrors(hipOccupancyMaxActiveBlocksPerMultiprocessor(
      &maxBlocksPerSM, pathcalc, blockSize, 0));

  const int gridSize = std::max(1, prop.multiProcessorCount * maxBlocksPerSM);
  const int totalThreads = gridSize * blockSize;

  std::printf("Using %d SMs, %d active blocks/SM, %d blocks total, %d threads\n",
              prop.multiProcessorCount, maxBlocksPerSM, gridSize, totalThreads);

  // allocate memory on host and device
  h_v = static_cast<float *>(std::malloc(sizeof(float) * NPATH));
  checkHipErrors(hipMalloc((void **)&d_v, sizeof(float) * NPATH));
  checkHipErrors(hipMalloc((void **)&state, sizeof(hiprandState_t) * totalThreads));

  // define constants and transfer to GPU
  h_T = 1.0f;
  h_r = 0.05f;
  h_sigma = 0.1f;
  h_rho = 0.5f;
  h_alpha = std::sqrt(1.0f - h_rho * h_rho);
  h_dt = 1.0f / h_N;
  h_con1 = 1.0f + h_r * h_dt;
  h_con2 = std::sqrt(h_dt) * h_sigma;

  checkHipErrors(hipMemcpyToSymbol(N, &h_N, sizeof(h_N)));
  checkHipErrors(hipMemcpyToSymbol(T, &h_T, sizeof(h_T)));
  checkHipErrors(hipMemcpyToSymbol(r, &h_r, sizeof(h_r)));
  checkHipErrors(hipMemcpyToSymbol(sigma, &h_sigma, sizeof(h_sigma)));
  checkHipErrors(hipMemcpyToSymbol(rho, &h_rho, sizeof(h_rho)));
  checkHipErrors(hipMemcpyToSymbol(alpha, &h_alpha, sizeof(h_alpha)));
  checkHipErrors(hipMemcpyToSymbol(dt, &h_dt, sizeof(h_dt)));
  checkHipErrors(hipMemcpyToSymbol(con1, &h_con1, sizeof(h_con1)));
  checkHipErrors(hipMemcpyToSymbol(con2, &h_con2, sizeof(h_con2)));

  // random number generator setup and kernel timing
  checkHipErrors(hipEventRecord(start, 0));
  setup_kernel<<<gridSize, blockSize>>>(state);
  getLastHipError("setup_kernel execution failed");
  pathcalc<<<gridSize, blockSize>>>(d_v, state, NPATH, totalThreads);
  getLastHipError("pathcalc execution failed");
  checkHipErrors(hipEventRecord(stop, 0));
  checkHipErrors(hipEventSynchronize(stop));
  checkHipErrors(hipEventElapsedTime(&milli, start, stop));

  std::printf("Monte Carlo kernel execution time (ms): %f\n", milli);

  // copy back results
  checkHipErrors(
      hipMemcpy(h_v, d_v, sizeof(float) * NPATH, hipMemcpyDeviceToHost));

  // compute average
  sum1 = 0.0;
  sum2 = 0.0;
  for (int i = 0; i < NPATH; i++) {
    sum1 += h_v[i];
    sum2 += h_v[i] * h_v[i];
  }

  std::printf("\n Average value and standard deviation of error = %13.8f %13.8f\n\n",
              sum1 / NPATH,
              std::sqrt((sum2 / NPATH - (sum1 / NPATH) * (sum1 / NPATH)) /
                        NPATH));

  // Release memory and exit cleanly
  std::free(h_v);
  checkHipErrors(hipFree(state));
  checkHipErrors(hipFree(d_v));

  // HIP exit -- needed to flush printf write buffer
  checkHipErrors(hipDeviceReset());
  return 0;
}