#include <R.h>
#include <Rinternals.h>
#include <Rdefines.h>
#include <cuda.h>
#include <math.h>

/* Illustrates calling 2-D kernels in R with C code on CPU and on GPU 
 Sergey V. 2022. */

extern "C" {
   SEXP kfunc_kernel_cpu(SEXP x1, SEXP x2, SEXP va, SEXP vb, SEXP vc, SEXP vkname);
   SEXP kfunc_kernel_gpu(SEXP x1, SEXP x2, SEXP va, SEXP vb, SEXP vc, SEXP vkname, SEXP devID);
}


void kfunc_exp_host(double *x1, double *x2, double *Sigma, int M, int N, double a, double b, double c){
 int i,j;
 for (i=0; i<M; i++) {
    for (j=0; j<N; j++){
      Sigma[j*M+i] = a*exp(-b*pow(fabs(x1[i]-x2[j]),c));
    }
  }
}


// CUDA exp kernel
__global__ void kfunc_exp_kernel(double *x1, double *x2, double *Sigma,  
                            const int M, 
                            const int N, 
                            const double a, const double b, const double c)
{
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  int j = threadIdx.y + blockIdx.y * blockDim.y;
  if(i<M && j<N){
      Sigma[j*M+i] = a*exp(-b*pow(fabs(x1[i]-x2[j]),c));
  }
  return;
}


void kfunc_matern_host(double *x1, double *x2, double *Sigma, int M, int N, double a, double b, double c){
 int i,j;
 for (i=0; i<M; i++) {
    for (j=0; j<N; j++){
      Sigma[j*M+i] = a*(1 + b*fabs(x1[i]-x2[j]))*exp(-c*fabs(x1[i]-x2[j]));
    }
  }
}



// Wrapper code between R and CPU based C 
SEXP kfunc_kernel_cpu(SEXP x1, SEXP x2, SEXP va, SEXP vb, SEXP vc, SEXP vkname){
   double *x1_host = NULL;
   double *x2_host = NULL;
   double *Sigma;
   double a = 0, b = 0, c = 0;
   int mm = 0, nn = 0;
   const char *kname;   
   SEXP Rdim;
 
   // data transfer from R to C 
   x1_host = REAL(x1);
   Rdim = getAttrib(x1, R_DimSymbol);
   mm = INTEGER(Rdim)[0];

   x2_host = REAL(x2);
   Rdim = getAttrib(x2, R_DimSymbol);
   nn = INTEGER(Rdim)[0];

   a = REAL(va)[0];
   b = REAL(vb)[0];
   c = REAL(vc)[0];
   kname = (char*)malloc(20*sizeof(char));
   kname = STRING_VALUE(vkname);

   printf("a,b,c = %4.3f, %4.3f, %4.3f\n", a,b,c);
   printf("%s\n", kname);

   Sigma = (double*)calloc(mm*nn,sizeof(double));

 
   // return value, allocated in C and can be used in R directly
   SEXP Rval;
   PROTECT(Rval = allocMatrix(REALSXP, mm, nn));

   if(strcmp(kname,"exponential")==0 || strcmp(kname,"exp")==0){
	printf("calling kfunc_exp..\n");
   	kfunc_exp_host(x1_host, x2_host, Sigma, mm, nn, a, b, c);
   } else if(strcmp(kname,"matern") == 0){
	printf("calling kfunc_matern..\n");
   	kfunc_matern_host(x1_host, x2_host, Sigma, mm, nn, a, b, c);
   }

   memcpy(REAL(Rval), Sigma, mm*nn*sizeof(double));

   free(Sigma);
   UNPROTECT(1);
   return Rval;
}


// Wrapper code between R and GPU based C 
SEXP kfunc_kernel_gpu(SEXP x1, SEXP x2, SEXP va, SEXP vb, SEXP vc, SEXP vkname, SEXP devID){
   double *x1_host = NULL;
   double *x1_d = NULL;
   double *x2_host = NULL;
   double *x2_d = NULL;
   double *Sigma;
   double *Sigma_d;
   double a = 0, b = 0, c = 0;
   int mm = 0, nn = 0, bs = 0;
   int gpuID = 0;
   const char *kname;   
   SEXP Rdim;
 
   // data transfer from R to C 
   x1_host = REAL(x1);
   Rdim = getAttrib(x1, R_DimSymbol);
   mm = INTEGER(Rdim)[0];

   x2_host = REAL(x2);
   Rdim = getAttrib(x2, R_DimSymbol);
   nn = INTEGER(Rdim)[0];

   a = REAL(va)[0];
   b = REAL(vb)[0];
   c = REAL(vc)[0];
   gpuID = INTEGER(devID)[0];
   kname = (char*)malloc(20*sizeof(char));
   kname = STRING_VALUE(vkname);

   printf("a,b,c = %4.3f, %4.3f, %4.3f\n", a,b,c);
   printf("%s\n", kname);
   printf("gpuID = %d\n", gpuID);

   Sigma = (double*)calloc(mm*nn,sizeof(double));

   // set GPU card num 
   cudaSetDevice(gpuID);
 
   // return value for R
   SEXP Rval;
   PROTECT(Rval = allocMatrix(REALSXP, mm, nn));

   // GPU memory allocation
   cudaMalloc(&x1_d, mm*sizeof(double));
   cudaMalloc(&x2_d, nn*sizeof(double));
   cudaMalloc(&Sigma_d, mm*nn*sizeof(double));

   if(Sigma_d == NULL) {
     printf("No space in GPU!");
     UNPROTECT(1);
     return R_NilValue;
   }

   // memory copy from CPU to GPU
   cudaMemcpy(x1_d, x1_host, mm*sizeof(double), cudaMemcpyHostToDevice); 
   cudaMemcpy(x2_d, x2_host, nn*sizeof(double), cudaMemcpyHostToDevice); 
   cudaMemcpy(Sigma_d, Sigma, mm*nn*sizeof(double), cudaMemcpyHostToDevice); 
 
   // CUDA kernel with blocks of threads
   bs = 32;
   dim3 block(32,32);
   dim3 grid( (mm*nn + bs - 1)/bs, (mm*nn + bs - 1)/bs );
   kfunc_exp_kernel<<<block,grid>>>(x1_d, x2_d, Sigma_d, mm, nn, a, b, c);
   
   //copy result from GPU
   cudaMemcpy(REAL(Rval), Sigma_d, mm*nn*sizeof(double), cudaMemcpyDeviceToHost); 
   cudaDeviceSynchronize();

   // Free GPU memory
   if(x1_d) {cudaFree(x1_d); x1_d=NULL;}
   if(x2_d) {cudaFree(x2_d); x2_d=NULL;}
   if(Sigma_d) {cudaFree(Sigma_d); Sigma_d=NULL;}
   UNPROTECT(1);
   return Rval;
}