#define MAX_GPU_COUNT 16

#include <fstream>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <cstdio>
#include <stdio.h>
#include <cstring>
#include <vector>
#include <cmath>
#include <omp.h>
#include <assert.h>
#include <cuda_runtime.h>

#include "device.hpp"
#include "string.hpp"
#include "help.hpp"
#include "lensFuncs.hpp"
#include "cudaLens.hpp"


using std::cout;
using std::endl;
using mycuda::toString;
using mycuda::Device;




int main(int argc, char* argv[]){
  //======================================== Verify Input and report errors ========================================
  params parameters;
  if ( setParams(argc,argv,parameters) ){//if successfull, setParams has already set the parameters.
    return -1;
  }
  //===================================================== END ======================================================

  
  //======================================== Parameter Initialization ========================================
  std::string target_dir       = parameters.target_dir;
  std::string lens_pos         = parameters.lens_pos;
  std::vector<int> device_list = parameters.devices;
  float  kappa_star         = parameters.kappa_star;
  float  kappa_s            = parameters.kappa_s;
  float  gamma              = parameters.gamma;
  int    resolution         = parameters.resolution;
  float  source_scale       = parameters.source_scale;
  float  avg_ray_count      = parameters.avg_ray_count;
  float  lens_scale_fudge_1 = parameters.lens_scale_fudge_1;
  float  lens_scale_fudge_2 = parameters.lens_scale_fudge_2;
  float  image_scale_fudge  = parameters.image_scale_fudge;
  float  mass               = parameters.mass;
  long   cycles             = parameters.cycles;
  long   cycles_done        = parameters.cycles_done;
  long   nray               = parameters.nray;
  long   exclude            = parameters.exclude;
  int    d_size             = device_list.size();

  // Calculate derived physical quantities
  float  kappa         = kappa_star + kappa_s;
  float  ray_grid_x    = (0.5 + 2.0*image_scale_fudge) * source_scale / (1.0 - kappa - gamma);
  float  ray_grid_y    = (0.5 + 2.0*image_scale_fudge) * source_scale / (1.0 - kappa + gamma);
  float  lens_rad_x    = (0.5*source_scale + lens_scale_fudge_1) / (1.0 - kappa + gamma);
  float  lens_rad_y    = (0.5*source_scale + lens_scale_fudge_1) / (1.0 - kappa - gamma);
  float  lens_rad      = sqrt(lens_rad_x*lens_rad_x + lens_rad_y*lens_rad_y) + lens_scale_fudge_2;
  long   lens_count    = (size_t)(kappa_star * lens_rad*lens_rad / mass + 0.5);
  float  ss2           = source_scale/2.0;
  // TODO: Eventually this should be the avg. lens mass
  // The rest of the code considers source_scale to be a radius, i.e., half the side length


  // Allocate and initialise lens data (determine own lens positions) on CPU.
  // Lenses x,y are important only to the 7th decimal, and m to the 2nd.
  ImagePlane IP;
  std::string temp = target_dir+"lens_pos.dat";
  char *lp_file    = (char*) temp.c_str();

  if ( lens_pos.compare("r") == 0 ) {
    srand48(timeMachineSeed(device_list));
    IP = setupLensPositions(lp_file,lens_count,lens_rad,lens_rad,mass);
  } else if ( atoi(lens_pos.c_str()) != 0 ) {
    srand48(atoi(lens_pos.c_str()));
    IP = setupLensPositions(lp_file,lens_count,lens_rad,lens_rad,mass);
  } else {
    std::ifstream in_stream(lp_file);
    if (in_stream){
      IP = readLensPositions(lp_file,&lens_count,lens_rad,lens_rad,mass);
    } else {
      cout << "Input file for lens positions does not exist." << endl;
      return -1;
    }
    in_stream.close();
  }


  // Allocate source plane array
  float** source_plane;
  source_plane = create2DArray<float>(resolution+1, resolution+1);



  cout << "gpu-d code version : 3.0"                     << endl;
  cout << "Output directory : "    << target_dir         << endl;
  cout << "****** Initialization ******"                 << endl;
  cout << " Kappa              = " << kappa              << endl;
  cout << " Kappa compact      = " << kappa_star         << endl;
  cout << " Kappa smooth       = " << kappa_s            << endl;
  cout << " Gamma              = " << gamma              << endl;
  cout << " Resolution         = " << resolution         << endl;
  cout << " Source scale       = " << source_scale       << endl;
  cout << " Lens Positions     = " << lens_pos           << endl;
  cout << " Lens count         = " << lens_count         << endl;
  cout << " Lens mass          = " << mass               << endl;
  cout << " Image scale x      = " << ray_grid_x*2.0     << endl;
  cout << " Image scale y      = " << ray_grid_y*2.0     << endl;
  cout << " Lens scale fudge 1 = " << lens_scale_fudge_1 << endl;
  cout << " Lens scale fudge 2 = " << lens_scale_fudge_2 << endl;
  cout << " Image scale fudge  = " << image_scale_fudge  << endl;
  cout << " Lens region rad x  = " << lens_rad_x         << endl;
  cout << " Lens region rad y  = " << lens_rad_y         << endl;
  cout << " Lens region radius = " << lens_rad           << endl;
  cout << " Avg. ray count     = " << avg_ray_count      << endl;
  printf("Doing %4li cycles on %1i devices with CHUNK_SIZE=%7i\n",cycles,d_size,CHUNK_SIZE);
  //================================================= END =================================================


  //======================================== Device Initialization ========================================
  std::vector<Device> devices;
  for( int d=0; d<d_size; ++d ) {
    Device device(device_list[d]);
    cout << "Using device " << toString(device.getIndex()) << " (" << device.getName() << ")";
    device.initialise();
    cout << " ... initialized" << endl;
    devices.push_back(device);
  }

  omp_set_num_threads(d_size); // IMPORTANT!!!
  //================================================= END =================================================


  //========================================== Memory Allocation ==========================================
  assert(IP.Nl < LENS_CHUNK);
  
  // Host memory pointers for ray positions
  float* x1[MAX_GPU_COUNT];
  float* x2[MAX_GPU_COUNT];
  float* y1[MAX_GPU_COUNT];
  float* y2[MAX_GPU_COUNT];
  // Device memory pointers for rays, lens positions and mass
  float* d_x1[MAX_GPU_COUNT];
  float* d_x2[MAX_GPU_COUNT];
  float* d_p1[MAX_GPU_COUNT];
  float* d_p2[MAX_GPU_COUNT];
  float* d_lx1[MAX_GPU_COUNT];
  float* d_lx2[MAX_GPU_COUNT];
  float* d_lm[MAX_GPU_COUNT];


  // Allocate Host memory for ray positions
  for( int d=0; d<d_size; ++d ) {
    x1[d] = (float*) calloc(CHUNK_SIZE, sizeof(float));
    x2[d] = (float*) calloc(CHUNK_SIZE, sizeof(float));
    y1[d] = (float*) calloc(CHUNK_SIZE, sizeof(float));
    y2[d] = (float*) calloc(CHUNK_SIZE, sizeof(float));
  }	

  // Allocate Device memory for rays, lens positions and mass (usage of Nleff for correct memory allocation)
#pragma omp parallel for
  for( int d=0; d<d_size; d++ ) {
    cudaMalloc( (void**) &d_x1[d], CHUNK_SIZE*sizeof(float));
    cudaMalloc( (void**) &d_x2[d], CHUNK_SIZE*sizeof(float));
    cudaMalloc( (void**) &d_p1[d], CHUNK_SIZE*sizeof(float));
    cudaMalloc( (void**) &d_p2[d], CHUNK_SIZE*sizeof(float));
    cudaMalloc( (void**) &d_lx1[d], IP.Nleff*sizeof(float));
    cudaMalloc( (void**) &d_lx2[d], IP.Nleff*sizeof(float));
    cudaMalloc( (void**) &d_lm[d],  IP.Nleff*sizeof(float));
      
    cudaError cuda_err = cudaGetLastError();
    if( cudaSuccess != cuda_err ){
      cout << "CUDA Error: Failed to allocate Device memory, device " << devices[d].getIndex() << endl;
    }
  }


  // Transfer lens positions and mass from Host to Device
#pragma omp parallel for
  for( int d=0; d<d_size; d++ ) {
    cudaMemcpy(d_lx1[d], IP.lx1, IP.Nl*sizeof(float),cudaMemcpyHostToDevice);
    cudaMemcpy(d_lx2[d], IP.lx2, IP.Nl*sizeof(float),cudaMemcpyHostToDevice);
    cudaMemcpy(d_lm[d],  IP.lm,  IP.Nl*sizeof(float),cudaMemcpyHostToDevice);
  
    cudaError cuda_err = cudaGetLastError();
    if( cudaSuccess != cuda_err ){
      cout << "CUDA Error: Failed to transfer lens position and mass Host memory to Device, device " << devices[d].getIndex() << endl;
    }
  }
  //================================================= END =================================================



  //======================================== Main loop ========================================
  cout << "****** Running ******" << endl;
 
  float last_out = cycles/10.;
  while ( last_out < (float)cycles_done ){
    last_out += (float)cycles/10.;
  }
  // Print initial status
  int pc      = (int)floor(100*(nray/CHUNK_SIZE)/cycles);
  float ex_pc = 0;
  float avgpp = (float)(nray-exclude)/(float)(resolution*resolution);
  printf("%3i%% complete, total rays: %12li, excluded: %12li (%7.4f%%), avg rays/pxl: %14.8f\n",pc,nray,exclude,ex_pc,avgpp);
  fflush(stdout);
  //if the following is done inside the loop then I keep allocating memory until I run out of it :
  float* yZeros = (float*) calloc(CHUNK_SIZE,sizeof(float));



  for( long m=cycles_done; m<cycles; m++ ) {
  //  for( long m=cycles_done; m<1; m++ ) {


    // Generate random starting ray positions
#pragma omp parallel for
    for( int d=0; d<d_size; d++ ) {
      genRandomPositions(x1[d], x2[d], CHUNK_SIZE, ray_grid_x, ray_grid_y, (long) m);
    }


    // Transfer starting ray positions from Host to Device
#pragma omp parallel for
    for( int d=0; d<d_size; d++ ) {
      cudaMemcpy(d_x1[d], x1[d], CHUNK_SIZE*sizeof(float), cudaMemcpyHostToDevice);
      cudaMemcpy(d_x2[d], x2[d], CHUNK_SIZE*sizeof(float), cudaMemcpyHostToDevice);
      //Overwrite the partial ys with zeros <- NECESSARY
      cudaMemcpy(d_p1[d], yZeros, CHUNK_SIZE*sizeof(float), cudaMemcpyHostToDevice);
      cudaMemcpy(d_p2[d], yZeros, CHUNK_SIZE*sizeof(float), cudaMemcpyHostToDevice);
      
      cudaError cuda_err = cudaGetLastError();
      if( cudaSuccess != cuda_err ){
	cout << "CUDA Error: Failed to transfer ray position Host memory to Device, device " << devices[d].getIndex() << endl;
      }
    }
    

    // Calculate deflections. Main calculation.
#pragma omp parallel for
    for( int d=0; d<d_size; d++ ) {
      runCudaLens(d_lx1[d], d_lx2[d], d_lm[d], d_x1[d], d_x2[d], d_p1[d], d_p2[d], IP.Nl, kappa_s, gamma);
    }


    // Transfer final ray positions from Device to Host
#pragma omp parallel for
    for( int d=0; d<d_size; d++ ) {
      cudaMemcpy(y1[d], d_x1[d], CHUNK_SIZE*sizeof(float), cudaMemcpyDeviceToHost);
      cudaMemcpy(y2[d], d_x2[d], CHUNK_SIZE*sizeof(float), cudaMemcpyDeviceToHost);

      cudaError cuda_err = cudaGetLastError();
      if( cudaSuccess != cuda_err ){
	cout << "CUDA Error: Failed to transfer final ray position Device memory to Host, device " << devices[d].getIndex() << endl;
      } 
    }


    // Bin the results in pixels
    for( int d=0; d<d_size; d++ ) {
      binSourcePlane(&nray, &exclude, source_plane, y1[d], y2[d], resolution, ss2);
    }

	
    // Print current status every 10%
    if( m >= last_out ) {
      last_out += (float)cycles/10.;
      //      writeSnapshot(target_dir,source_plane,kappa,gamma,ss2,ray_grid_x,ray_grid_y,nray,exclude,resolution);
      reportProgress(cycles,nray,exclude,resolution);
    }


  }


  // Print out the final ray statistics
  writeSnapshot(target_dir,source_plane,kappa,gamma,ss2,ray_grid_x,ray_grid_y,nray,exclude,resolution);
  reportProgress(cycles,nray,exclude,resolution);
  //================================================= END =======================================================


  //=========================================== Cleanup GPU =====================================================
  // Loop over devices
#pragma omp parallel for
  for(  int d=0; d<d_size; ++d ) {
    cudaFree(d_x1[d]);
    cudaFree(d_x2[d]);
    cudaFree(d_p1[d]);
    cudaFree(d_p2[d]);
  }
  //=============================================== END =========================================================


  cout << "Done." << endl;


  return 0;
}
    
