#include <iostream>
#include <fstream>
#include <exception>
#include <string>
#include <sstream>
#include <set>
#include <vector>
#include <iterator>
using namespace std;

#include <cuda.h>
#include <stdio.h>

// Declare our texture reference. This MUST be a global variable!
//                  float2 = texel element is 2-D value using CUDA array type float2
//       cudaTextureType2D = indexing will use 2-dimensions
// cudaReadModeElementType = Values returned will be same type as stored (no
//                           transformation to normalized values)
static texture<float2, cudaTextureType2D, cudaReadModeElementType> bfield_GPU_texref;

float2* ReadBfield(int &Nzbins, float &zmin, float &zmax, int &Nrbins, float &rmin, float &rmax);

//......................
// SimpleException
//......................
class SimpleException:public std::exception
{
 public:
  SimpleException(string str){ msg = str;}
  virtual ~SimpleException() throw(){}
  virtual const char* what(void) const throw() { return msg.c_str(); }
 private:
  string msg;
};

//----------------
// InitCUDA
//----------------
void InitCUDA(bool quiet=false)
{
  // Verify that a CUDA device exists and optionally print some
  // info about it.

  int deviceCount=0;
  cudaGetDeviceCount(&deviceCount);
  if(deviceCount<1){
    cout<<"No CUDA devices available!"<<endl;
    throw SimpleException("No CUDA devices available!");
  }

  // Loop over devices, listing them all
  cout << endl;
  cout << "CUDA GPU devices found: "<<deviceCount<<endl;
  cout << "---------------------"<<endl;
  for(int device=0; device<deviceCount; device++){
  
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, device);

    // Form useful string for properties
    stringstream ss;
    ss << device << ".) ";
    ss << prop.name << " ";
    ss << prop.totalGlobalMem/1024/1024/1024 << "GB ";
    ss << (int)(prop.clockRate/1000.0) << "MHz "; // clock rate is reported in kHz by device
    ss << "CUDA " << prop.major << "." << prop.minor << " ";

    cout << ss.str() << endl;
  }
  cout << endl;
  
  int my_device;
  cudaGetDevice(&my_device);
  cout << "Using CUDA GPU device "<< my_device <<endl;
  cout << "---------------------"<<endl;

}

//----------------
// CopyBfieldMapToGPU
//----------------
void CopyBfieldMapToGPU(void)
{
  cudaError_t err;

  // Read in Bfield to host memory
  cout << "Reading in B-field map..." << endl;
  int Nzbins, Nrbins;
  float zmin, zmax;
  float rmin, rmax;
  float2* bfield_table = ReadBfield(Nzbins, zmin, zmax, Nrbins, rmin, rmax);
  cudaExtent extent =  make_cudaExtent(Nzbins, Nrbins, 0);

  // Allocate a section of linear pitched memory on the device and
  // copy the map into it.
  float2 *d_bfield_table;
  size_t pitch;
  size_t rowsize = sizeof(float2)*Nzbins;
  err = cudaMallocPitch(&d_bfield_table, &pitch, rowsize, (size_t)Nrbins);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));
  err = cudaMemcpy2D(d_bfield_table, pitch, bfield_table, rowsize, rowsize, Nrbins, cudaMemcpyHostToDevice);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));

  // Inform user how many bytes of global memory we are using
  //-----------------------------------------------------------------
  int memsize_kB = pitch*Nrbins/1024;
  cout << "Allocated memory on GPU device ("<<memsize_kB<<"kB)"<<endl;

  // Configure texture object
  bfield_GPU_texref.addressMode[0] = bfield_GPU_texref.addressMode[1] = cudaAddressModeBorder;
  bfield_GPU_texref.filterMode = cudaFilterModeLinear;
  bfield_GPU_texref.normalized = true;

  // Create a channel descriptor that says the data will be stored as float2
  cudaChannelFormatDesc desc = cudaCreateChannelDesc<float2>();

  // Bind texture object to allocated memory on device
  cout << "Binding B-field to texture object ...." << endl;
  err = cudaBindTexture2D(NULL, &bfield_GPU_texref, d_bfield_table, &desc, Nzbins, Nrbins, pitch);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));

#if 0
  // This commented out section tries to allocate and copy the map
  // as a cuda3D array object which in principle, is supposed to be
  // a better way of doing it (saving us from having to worry about
  // pitched memory for one thing). However, it never worked and most
  // of the web documentation centered on pitched pointer allocation
  // so I reverted to that. Leaving this here temporarily in case
  // additional inspiration arises to have another run at this method.

  // Create a channel descriptor that says the data will be stored as float2
  cudaChannelFormatDesc desc = cudaCreateChannelDesc<float2>();
  // Allocate the array on the GPU device
  cudaArray_t bfield_GPU;
  err = cudaMalloc3DArray(&bfield_GPU, &desc, extent);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));

  // A brief note on CUDA pitched pointers:
  // A CUDA pitched pointer is used to describe the format of the
  // data in the memory allocated on the host. The "pitch" is how
  // many bytes are allocated for each row and the xsz and ysz
  // parameters tell how many elements per row and rows in the 
  // table repectively. This system allows one to ensure that the
  // rows are properly byte aligned on the host (in case there is
  // some need for that). We use the "make_cudaPitchedPtr" routine
  // to do this below.

  // Copy the Bfield from host memory to device memory
  cout << "Copying Bfield map to GPU ..." << endl;
  cudaMemcpy3DParms myparms = {0};
  myparms.srcPtr = make_cudaPitchedPtr((void*)bfield_table, sizeof(float2)*Nzbins, Nzbins, Nrbins);
  myparms.dstArray = bfield_GPU;
  myparms.extent = extent;
  myparms.kind = cudaMemcpyHostToDevice;
  err = cudaMemcpy3D(&myparms);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));

  // Create texture object
  bfield_GPU_texref.addressMode[0] = bfield_GPU_texref.addressMode[1] = cudaAddressModeBorder;
  //bfield_GPU_texref.filterMode = cudaFilterModeLinear;
  bfield_GPU_texref.filterMode = cudaFilterModePoint;
  bfield_GPU_texref.normalized = true;

  // Bind texture object to allocated array on device
  cout << "Binding B-field to texture object ...." << endl;
  err = cudaBindTextureToArray(&bfield_GPU_texref, bfield_GPU, &desc);
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));
#endif

  cout << "B-field ready on GPU." << endl;
}

//----------------
// ReadBfield
//----------------
float2* ReadBfield(int &Nzbins, float &zmin, float &zmax, int &Nrbins, float &rmin, float &rmax)
{
  // Open the input file and read in all values to figure out how
  // many bins we have in z and r.
  string fname = "solenoid_1500_poisson_20090814_01";
  ifstream ifs(fname.c_str());
  if(!ifs.is_open()) throw SimpleException(string("Unable to open B-field map file: " + fname));

  // Read all elements into a vector, storing x(=r) and z coordinates in
  // separate set containers so we can know how many elements and what their
  // range is for each dimension.
  set<float> z_vals;
  set<float> r_vals;
  vector<float4> map_vals; // .x=r .y=z .z=Br .w=Bz
  while(ifs.good()){
    char line[1024];
    ifs.getline(line, 1024);
    if(!ifs.good()) break;
    if(line[0] == '#') continue;
    stringstream ss(line);
    float x,y,z, Bx, By, Bz;
    ss>>x>>y>>z>>Bx>>By>>Bz;
    x *= 2.54; // convert to inches
    z *= 2.54; // convert to inches
    z_vals.insert(z);
    r_vals.insert(x);
    float4 tmp = {x, z, Bx, Bz};
    map_vals.push_back(tmp);
  }
  
  // Copy ranges to caller's variables and report map stats to user
  Nzbins = z_vals.size();
  Nrbins = r_vals.size();
  rmin = *(r_vals.begin());  // assume set orders by ascending value!
  rmax = *(r_vals.rbegin()); // assume set orders by ascending value!
  zmin = *(z_vals.begin());  // assume set orders by ascending value!
  zmax = *(z_vals.rbegin()); // assume set orders by ascending value!
  cout << "  Found "<<Nzbins<<" z-bins and "<<Nrbins<<" r-bins" << endl;
  cout << "    R-range:"<<rmin<<" to "<<rmax<<endl;
  cout << "    Z-range:"<<zmin<<" to "<<zmax<<endl;

  // Allocate memory to hold table on host side in a format that 
  // is compatible with the eventual texture map on the device side.
  float2 *bfield_table = (float2*)malloc(sizeof(float2) * Nzbins * Nrbins);
  
  // Copy the map into the allocated memory
  set<int> indexes;
  for(unsigned int i=0; i<map_vals.size(); i++){
    float r = map_vals[i].x;
    float z = map_vals[i].y;
    float Br = map_vals[i].z;
    float Bz = map_vals[i].w;

    int rbin = distance(r_vals.begin(), r_vals.find(r));
    int zbin = distance(z_vals.begin(), z_vals.find(z));
    int index = zbin + Nzbins*rbin;
    indexes.insert(index);

    // In the texture map, the "x" element of the float2 will be Bz
    // and the "y" element will be Br
    // Bz=4.56;
    // Br=6.54;
    bfield_table[index].x = Bz;
    bfield_table[index].y = Br;
  }
  cout <<"    " << indexes.size() << " elements entered ranging from " << *(indexes.begin()) << " to " << *(indexes.rbegin()) << endl;

  return bfield_table;
}

//----------------
// GPU_GetBfield
//----------------
__global__ void GPU_GetBfield(int N, float2 *d_coordinates, float2 *d_B)
{
  int index = threadIdx.x + blockDim.x*blockIdx.x;
  if(index<N){
    float z = d_coordinates[index].x;
    float r = d_coordinates[index].y;

    d_B[index] = tex2D(bfield_GPU_texref, z, r);
  }

}

//----------------
// TestGPUBfieldMap
//----------------
void TestGPUBfieldMap(void)
{
  cout << endl;
  cout << "Testing B-field map reads via GPU texture memory ..." <<endl;

  // Generate some random points to evaluate the B-field at
  vector<float2> points;
  for(int i=0; i<10; i++){
    float z = -100.0 + 700.0*(float)random()/(float)RAND_MAX;
    float r = 250.0*(float)random()/(float)RAND_MAX;
    z=(float)random()/(float)RAND_MAX;
    r=(float)random()/(float)RAND_MAX/4.0;
    float2 tmp = {z,r};
    points.push_back(tmp);
  }
  cout << "   " << points.size() << " points generated for testing" <<endl;

  // Allocate some memory on device and fill it with coordinates
  float2 *d_coordinates, *d_B;
  int Nbytes = sizeof(float2)*points.size();
  cudaMalloc(&d_coordinates, Nbytes);
  cudaMalloc(&d_B, Nbytes);
  cudaMemcpy(d_coordinates, &points[0], Nbytes, cudaMemcpyHostToDevice);

  // Launch kernel to evaluate field
  cout << "   Launching GPU kernel ..." << endl; 
  GPU_GetBfield<<<1,10>>>(points.size(), d_coordinates, d_B);

  // Copy memory back to host
  float2 *h_B = (float2*)malloc(Nbytes);
  cudaMemcpy(h_B, d_B, Nbytes, cudaMemcpyDeviceToHost);
  
  // Print results to screen
  cout << "   Results:" << endl;
  cout << "   ---------------" << endl;
  for(int i=0; i<10; i++){
    float z = -100.0 + 700.0* points[i].x;
    float r =    0.0 + 250.0* points[i].y;
    float Bz =  h_B[i].x;
    float Br = h_B[i].y;
    cout << "     (z,r)=(" << z << ", " << r << ")";
    cout << "   (Bz,Br)=(" << Bz << ", " << Br << ")";
    cout << endl;
  }

  cout << "   Done." <<endl;
}


//----------------
// main
//----------------
int main(int narg, char *argv[])
{
  InitCUDA();
  CopyBfieldMapToGPU();
  TestGPUBfieldMap();

  return 0;
}