/*============================================================================

 Function      vfiGPU

 Usage         vfiGPU(hV, hG)

 Arguments     hV: pointer to array of doubles storing the value function.
                   
               hG: pointer to array of doubles storing the policy function.
	              
 Description   This function performs value function iteration on the GPU,
               finding the maximum of the Bellman objective function for each
	       node in the state space and iterating until convergence.

 Dependencies  Global Variables: eta, beta, alpha, delta, mu, rho, sigma,
                                 block_size, nk, nz, tol, maxtype, howard
				 (globalvas.h).

               Functions:        pow (math.h);
	                         cblasDaxpy, cblasIdamax (cblas.h).

	       Kernels:          ar1 (ar1.cu), kGrid (kGrid.cu),
	                         vfInit(vfInit.cu), vfStep (vfStep.cu).

 Return value  Returns 0 upon successful completion, 1 otherwise.

 =============================================================================

 Author:       Eric M. Aldrich

 Contact:      ealdrich@gmail.com

 Date:         20 April 2010

 ============================================================================*/

#include "globalvars.h"
#include "cublas.h"
#include <iostream>
#include <ctime>

using namespace std;

#include "ncdf.cu"
#include "binary_val.cu"
#include "ar1.cu"
#include "kGrid.cu"
#include "vfInit.cu"
#include "grid_max.cu"
#include "binary_max.cu"
#include "vfStep.cu"

int vfiGPU(double* hV, double* hG) 
{ 

  int imax;
  double diff = 1.0;

  // pointers to variables in device memory
  double* K;
  double* Z;
  double* P;
  double* V0;
  double* V;
  double* G;
  double* Vtemp;

  // allocate variables in device memory
  size_t sizeK = nk*sizeof(double);
  size_t sizeZ = nz*sizeof(double);
  size_t sizeP = nz*nz*sizeof(double);
  size_t sizeV = nk*nz*sizeof(double);
  size_t sizeG = nk*nz*sizeof(double);
  clock_t start = clock();
  cudaMalloc((void**)&K, sizeK);
  cudaMalloc((void**)&Z, sizeZ);
  cudaMalloc((void**)&P, sizeP);
  cudaMalloc((void**)&V0, sizeV);
  cudaMalloc((void**)&Vtemp, sizeV);
  cudaMalloc((void**)&V, sizeV);
  cudaMalloc((void**)&G, sizeG);
  cout << "GPU Memory Allocation: " << (clock() - start)/(double)CLOCKS_PER_SEC << endl;

  // blocking
  dim3 dimBlockZ(nz, 1);
  dim3 dimBlockK(block_size,1);
  dim3 dimBlockV(block_size, nz);
  dim3 dimGridZ(1,1);
  dim3 dimGridK(nk/block_size,1);
  dim3 dimGridV(nk/block_size,1);

  // compute TFP grid, capital grid and initial VF
  double lambda = 3;
  ar1<<<dimGridZ,dimBlockZ>>>(nz,lambda,mu,sigma,rho,Z,P);
  kGrid<<<dimGridK,dimBlockK>>>(nk,nz,beta,alpha,delta,Z,K);
  vfInit<<<dimGridV,dimBlockV>>>(nz,eta,beta,alpha,delta,Z,V0);

  // iterate on the value function
  int count = 0;
  bool how = false;
  start = clock();
  while(fabs(diff) > tol){
    if(count < 3 | count % howard == 0) how = false; else how = true;
    vfStep<<<dimGridV,dimBlockV>>>(nk,nz,eta,beta,alpha,delta,maxtype,how,K,Z,P,V0,V,G);
    cublasDaxpy(nk*nz, -1.0, V, 1, V0, 1); 
    imax = cublasIdamax(nk*nz, V0, 1);
    cudaMemcpy(&diff, V0+imax, sizeof(double), cudaMemcpyDeviceToHost);
    Vtemp = V0;
    V0 = V;
    V = Vtemp;
    ++count;
  }
  cout << "GPU Solve Time: " << (clock() - start)/(double)CLOCKS_PER_SEC << endl;
  V = V0;
  
  // copy value and policy functions to host memory
  cudaMemcpy(hV, V, sizeV, cudaMemcpyDeviceToHost);
  cudaMemcpy(hG, G, sizeG, cudaMemcpyDeviceToHost);

  // copy state variable grids and transition matrix to host memory
  double* hK = new double[nk];
  double* hZ = new double[nz];
  double* hP = new double[nz*nz];
  cudaMemcpy(hK, K, sizeK, cudaMemcpyDeviceToHost);
  cudaMemcpy(hZ, Z, sizeZ, cudaMemcpyDeviceToHost);
  cudaMemcpy(hP, P, sizeP, cudaMemcpyDeviceToHost);

  // free variables in device memory
  cudaFree(K);
  cudaFree(Z);
  cudaFree(P);
  cudaFree(V0);
  cudaFree(V);
  cudaFree(Vtemp);
  cudaFree(G);

  return 0;

}
