#ifndef _GPUMAGNETICFIELD_CU_
#define _GPUMAGNETICFIELD_CU_

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

#include "SimpleException.h"
#include "GPUMagneticField.h"


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

/// This global variable is declared in device constant memory
/// to hold the map deminsions info so that it can be accessed quickly
/// and without having to pass a reference to the structure with every call.
__device__ __constant__ Bfield_parms_t BFIELD_PARMS;


// Local routines not exposed as part of the library
static float4* ReadBfieldFromFile(string fname, Bfield_parms_t &bparms);


//----------------
// InitCUDA
//----------------
void InitCUDA(bool quiet)
{
  /// Verify that a CUDA device exists and optionally
  /// print some info about it. This really does nothing
  /// to set up the device, just verifies one exists
  /// and informs the user what type of GPUs exist and
  /// which one we are using. Calling InitCUDA is not
  /// required.

  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(string fname)
{
  /// This is used to read in a specified B-field map
  /// (in ascii format with "x y z Bx By Bz" values
  /// on each line.) It will then set up a CUDA
  /// texture reference to point to the map so it
  /// can be accessed using the GPU_GetBfield*
  /// routines.

  cudaError_t err;

  // Read in Bfield to host memory
  cout << "Reading in B-field map..." << endl;
  Bfield_parms_t bparms;
  float4* bfield_table = ReadBfieldFromFile(fname, bparms);
  int &Nzbins = bparms.Nzbins;
  int &Nrbins = bparms.Nrbins;
  cudaExtent extent =  make_cudaExtent(Nzbins, Nrbins, 0);

  // Copy B-field map parameters to the device constant memory
  err = cudaMemcpyToSymbol(BFIELD_PARMS, &bparms, sizeof(Bfield_parms_t));
  if(err != cudaSuccess) throw SimpleException(string("CUDA ERROR: ") + cudaGetErrorString(err));

  // Allocate a section of linear pitched memory on the device and
  // copy the map into it.
  float4 *d_bfield_table;
  size_t pitch;
  size_t rowsize = sizeof(float4)*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] = cudaAddressModeClamp;
  bfield_GPU_texref.filterMode = cudaFilterModeLinear;
  bfield_GPU_texref.normalized = true;

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

  // 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));

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

//----------------
// ReadBfield
//----------------
float4* ReadBfieldFromFile(string fname, Bfield_parms_t &bparms)
{
  /// This opens an ASCII file containing the b-field map formatted
  /// to have the values "x y z Bx By Bz" on each line. The "x"
  /// and "Bx" values are taken to be "r" and "Br" while both
  /// of the y values are ignored. The number of bins and range
  /// are determined by the entries found in the file. Values
  /// are assumed to be in inches and Tesla.

  // 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 cm from inches
    z *= 2.54; // convert to cm from 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
  bparms.Nzbins = z_vals.size();
  bparms.Nrbins = r_vals.size();
  bparms.rmin = *(r_vals.begin());  // assume set orders by ascending value!
  bparms.rmax = *(r_vals.rbegin()); // assume set orders by ascending value!
  bparms.zmin = *(z_vals.begin());  // assume set orders by ascending value!
  bparms.zmax = *(z_vals.rbegin()); // assume set orders by ascending value!
  cout << "  Found "<<bparms.Nzbins<<" z-bins and "<<bparms.Nrbins<<" r-bins" << endl;
  cout << "    R-range:"<<bparms.rmin<<" to "<<bparms.rmax<<endl;
  cout << "    Z-range:"<<bparms.zmin<<" to "<<bparms.zmax<<endl;

  // It is quicker to multiply by a number than divide by a difference
  // so pre-calculate the scale factors for r and z coordinates
  bparms.r_scale = 1.0/(bparms.rmax - bparms.rmin);
  bparms.z_scale = 1.0/(bparms.zmax - bparms.zmin);

  // Allocate memory to hold table on host side in a format that 
  // is compatible with the eventual texture map on the device side.
  float4 *bfield_table = (float4*)malloc(sizeof(float4) * bparms.Nzbins * bparms.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;
    float B = sqrt(Br*Br + Bz*Bz);

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

    // In the texture map:
    //  "x" = normalized Br
    //  "y" = <unused>
    //  "z" = normalized Bz
    //  "w" = field magnitude
    bfield_table[index].x = (B==0.0 ? 0.0:Br/B);
    bfield_table[index].y = 0.0;
    bfield_table[index].z = (B==0.0 ? 0.0:Bz/B);
    bfield_table[index].w = B;
  }
  cout <<"    " << indexes.size() << " elements entered ranging from " << *(indexes.begin()) << " to " << *(indexes.rbegin()) << endl;

  return bfield_table;
}

//----------------
// GPU_GetBfieldPoints
//----------------
__global__ void GPU_GetBfieldPoints(int N, float3 *d_coordinates, float3 *d_B)
{
  /// This is called from the host to evaluate the
  /// field at a set of points. It is primarily
  /// intended for testing.

  int index = threadIdx.x + blockDim.x*blockIdx.x;
  if(index<N){
    GPU_GetBfield(d_coordinates[index], &d_B[index]);
  }
}

//----------------
// GPU_GetBfield
//----------------
static __device__ void GPU_GetBfield(float3 pos, float3 *B)
{
  /// This is a convenience method that just wraps the
  /// GPU_GetBfieldXYZ routine so the field can be
  /// accessed using float3 instead of individual float
  /// values.

  GPU_GetBfieldXYZ(pos.x, pos.y, pos.z, &B->x, &B->y, &B->z);
}

//----------------
// GPU_GetBfieldXYZ
//----------------
static __device__ void GPU_GetBfieldXYZ(float x, float y, float z, float *Bx, float *By, float *Bz)
{
  /// Access the field map in order to return the value
  /// of the B-field at the specified location. The
  /// field map is kept in r and z only so rotation
  /// into and out of that must be done for every call.
  /// If the r and z values in lab coordinates are already
  /// known, the GPU_GetBfieldRZ routine will be faster.
  /// The map is stored as magnitude and direction components
  /// so it will be faster to use the other GPU_GetBfieldXYZ
  /// routine, especially if you need the magnitude since 
  /// it will save calculating it from the components.

  // convert into cylindrical coordinates
  float phi = atan2(y, x);
  float r = sqrt(x*x + y*y);

  // calculate normalized coordinates
  float znorm = (z - BFIELD_PARMS.zmin)*BFIELD_PARMS.z_scale;
  float rnorm = (r - BFIELD_PARMS.rmin)*BFIELD_PARMS.r_scale;

  // evaluate map at specified coordinates
  float4 tmp = tex2D(bfield_GPU_texref, znorm, rnorm);

  // copy results into caller-supplied vector,
  // converting back into 3D coordinates
  float B = tmp.w;
  float Br = B*tmp.x;
  *Bx = Br*cos(phi);
  *By = Br*sin(phi);
  *Bz = B*tmp.z;
}

//----------------
// GPU_GetBfieldRZ
//----------------
static __device__ void GPU_GetBfieldRZ(float r, float z, float *Br, float *Bz)
{
  /// Access the field map in order to return the value
  /// of the B-field at the specified location in cylindrical
  /// coordinates.
  /// The map is stored as magnitude and direction components
  /// so it will be faster to use the other GPU_GetBfieldRZ
  /// routine, especially if you need the magnitude since 
  /// it will save calculating it from the components.

  // calculate normalized coordinates
  float znorm = (z - BFIELD_PARMS.zmin)*BFIELD_PARMS.z_scale;
  float rnorm = (r - BFIELD_PARMS.rmin)*BFIELD_PARMS.r_scale;

  // evaluate map at specified coordinates
  float4 tmp = tex2D(bfield_GPU_texref, znorm, rnorm);

  // copy results into caller-supplied vector,
  // converting back into 3D coordinates
  float B = tmp.w;
  *Br = B*tmp.x;
  *Bz = B*tmp.z;
}

//----------------
// GPU_GetBfield
//----------------
static __device__ void GPU_GetBfield(float3 pos, float4 *B)
{
  /// This is a convenience method that just wraps the
  /// GPU_GetBfieldXYZ routine so the field can be
  /// accessed using float3 and float 4 instead of
  /// individual float values.

  GPU_GetBfieldXYZ(pos.x, pos.y, pos.z, &B->w, &B->x, &B->y, &B->z);
}

//----------------
// GPU_GetBfieldXYZ
//----------------
static __device__ void GPU_GetBfieldXYZ(float x, float y, float z, float *B, float *bx, float *by, float *bz)
{
  /// Access the field map in order to return the value
  /// of the B-field at the specified location. The
  /// field map is kept in r and z only so rotation
  /// into and out of that must be done for every call.
  /// If the r and z values in lab coordinates are already
  /// known, the GPU_GetBfieldRZ routine will be faster.
  /// The map is stored as magnitude and direction components
  /// so it will be faster to use the other GPU_GetBfieldXYZ
  /// routine, especially if you need the magnitude since 
  /// it will save calculating it from the components.

  // convert into cylindrical coordinates
  float phi = atan2(y, x);
  float r = sqrt(x*x + y*y);

  // calculate normalized coordinates
  float znorm = (z - BFIELD_PARMS.zmin)*BFIELD_PARMS.z_scale;
  float rnorm = (r - BFIELD_PARMS.rmin)*BFIELD_PARMS.r_scale;

  // evaluate map at specified coordinates
  float4 tmp = tex2D(bfield_GPU_texref, znorm, rnorm);

  // copy results into caller-supplied vector,
  // converting back into 3D coordinates
  *B  = tmp.w;
  float &br = tmp.x;
  *bx = br*cos(phi);
  *by = br*sin(phi);
  *bz = tmp.z;
}

//----------------
// GPU_GetBfieldRZ
//----------------
static __device__ void GPU_GetBfieldRZ(float r, float z, float *B, float *br, float *bz)
{
  /// Access the field map in order to return the value
  /// of the B-field at the specified location in cylindrical
  /// coordinates.
  /// The map is stored as magnitude and direction components
  /// so it will be faster to use the other GPU_GetBfieldRZ
  /// routine, especially if you need the magnitude since 
  /// it will save calculating it from the components.

  // calculate normalized coordinates
  float znorm = (z - BFIELD_PARMS.zmin)*BFIELD_PARMS.z_scale;
  float rnorm = (r - BFIELD_PARMS.rmin)*BFIELD_PARMS.r_scale;

  // evaluate map at specified coordinates
  float4 tmp = tex2D(bfield_GPU_texref, znorm, rnorm);

  // copy results into caller-supplied vector,
  // converting back into 3D coordinates
  *B  = tmp.w;
  *br = tmp.x;
  *bz = tmp.z;
}

#endif // _GPUMAGNETICFIELD_CU_