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


#include "GPUMagneticFieldStepper.cu"

//----------------
// 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<float3> points;
  for(int i=0; i<10; i++){
    float x = 80.0*(float)random()/(float)RAND_MAX;
    float y = 80.0*(float)random()/(float)RAND_MAX;
    float z = -100.0 + 700.0*(float)random()/(float)RAND_MAX;
    float3 tmp = {x,y,z};
    points.push_back(tmp);
  }
  cout << "   " << points.size() << " points generated for testing" <<endl;

  // Allocate some memory on device and fill it with coordinates
  float3 *d_coordinates, *d_B;
  int Nbytes = sizeof(float3)*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_GetBfieldPoints<<<1,10>>>(points.size(), d_coordinates, d_B);

  // Copy memory back to host
  float3 *h_B = (float3*)malloc(Nbytes);
  cudaMemcpy(h_B, d_B, Nbytes, cudaMemcpyDeviceToHost);
  
  // Print results to screen
  cout << "   Results:" << endl;
  cout << "   ---------------" << endl;
  cout<<setprecision(3);
  for(int i=0; i<10; i++){
    float x = points[i].x;
    float y = points[i].y;
    float z = points[i].z;
    float Bx =  h_B[i].x;
    float By =  h_B[i].y;
    float Bz =  h_B[i].z;
    cout << "     (x,y,z)=(" << x << ", " << y << ", " << z << ")";
    cout << "  (Bx,By,Bz)=(" << Bx << ", " << By << ", " << Bz << ")";
    cout << endl;
  }

  cudaFree(d_coordinates);
  cudaFree(d_B);

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

//----------------
// TestGPUSwimmer
//----------------
void TestGPUSwimmer(void)
{
	// Define starting poisiton and momentum
	float3 pos = make_float3(0.0, 0.0, 65.0); // start at center of target
	float3 mom = make_float3(0.350, 0.0, 1.000); // 1.06GeV going at ~19.2 degrees
	float q = 1.0;
	float mass = 0.140; // pion mass (unused at the moment)
	trajectory_t *traj = NULL; // have GPU_Swim allocate memory for us
	
	Swim(q, mass, pos, mom, traj);

	state_t *step = traj->steps;
	for(int i=0; i<10; i++, step++){

		if(i>=traj->Nsteps)break;
		
		float3 &pos = step->pos;
		float3 &mom = step->mom;
		float epsilon = 1.0E-3; // this corresponds roughly to 1% resolution of a 12GeV particle	
		float p = 1.0/(fabs(step->q_over_p)+epsilon); // epsilon handles case when q_over_p is zero

		mom.x *= p;
		mom.y *= p;
		mom.z *= p;

		cout << "(x,y,z)=("<<pos.x<<", "<<pos.y<<", "<<pos.z<<")  ";
		cout << "(px,py,pz)=("<<mom.x<<", "<<mom.y<<", "<<mom.z<<")  ";
		cout << "mass="<<step->mass<< "  p="<< p << endl;
	}
 
}

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

 	return 0;
}