#include <iostream>
using namespace std;

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

#include "prof_time.h"
// This example tests memory access speeds using a B-field map

#define BZ_MIN -50.0
#define BZ_MAX 650.0
#define BZ_N 700

#define BR_MIN 0.0
#define BR_MAX 100.0
#define BR_N 100

typedef struct{
	float r, z;
	float Br, Bz;
	float dBrdr, dBzdr;
	float dBrdz, dBzdz;
}Bfield_Map_t;


void TestAccessSpeedOnHost(Bfield_Map_t* bfield, int Naccesses, Bfield_Map_t *checksums);
__global__ void TestAccessSpeedOnDevice(Bfield_Map_t* bfield, int Naccesses, Bfield_Map_t *checksums);


//--------------------------
// main
//--------------------------
int main(int narg, char *argv[])
{
	// Start hi-res timers
	prof_time::StartTimers();
	map<string, prof_time::time_diffs> prof_times;

	// Get number of CUDA devices to make sure one is available to run on
	int deviceCount=0;
	cudaGetDeviceCount(&deviceCount);
	if(deviceCount<1){
		cout<<"No CUDA devices available!"<<endl;
		return -1;
	}
	cout<<"There are "<<deviceCount<<" CUDA devices available."<<endl;

	// Allocate memory for B-field map on host
	size_t size = BZ_N*BR_N*sizeof(Bfield_Map_t);
	cout<<"Allocating "<<size/1024/1024<<"MB for B-field map"<<endl;
	Bfield_Map_t *bfield_h = (Bfield_Map_t*)malloc(size);
	
	// Alocate on device
	Bfield_Map_t *bfield_d;
	cudaError_t err = cudaMalloc(&bfield_d, size);
	cout<<"Allocating memory on device: "<<cudaGetErrorString(err)<<endl;
	
	// Fill in map values
	cout<<"Filling in map values"<<endl;
	float delta_r = (BR_MAX-BR_MIN)/(float)(BR_N-1);
	float delta_z = (BZ_MAX-BZ_MIN)/(float)(BZ_N-1);
	for(int ir=0; ir<BR_N; ir++){
		float r = BR_MIN + (float)ir*delta_r;
		for(int iz=0; iz<BZ_N; iz++){
			float z = BZ_MIN + (float)iz*delta_z;
			Bfield_Map_t *B = &bfield_h[iz*BR_N + ir];
			
			B->r = r;
			B->z = z;
			B->Br = sin(3.14159265*r/100.0);
			B->Bz = -2.0 + 0.2*cos(3.14159265*z/200.0);
			B->dBrdr = -cos(3.14159265*r/100.0);
			B->dBzdr = 0.0;
			B->dBrdz = 0.0;
			B->dBzdz = 0.2*sin(3.14159265*z/200.0);			
		}
	}
	
	// Copy map to device
	cout<<"Copying map to device ..."<<endl;
	cudaMemcpy(bfield_d, bfield_h, size, cudaMemcpyHostToDevice);
	
	// Test access speed on host
	int Naccesses = 1000000;
	cout<<"Testing "<<Naccesses/1000000<<"M accesses on host"<<endl;
	Bfield_Map_t *checksums_h = (Bfield_Map_t*)malloc(sizeof(Bfield_Map_t));
	prof_time start_time_host;
	TestAccessSpeedOnHost(bfield_h, Naccesses, checksums_h);
	start_time_host.TimeDiffNow(prof_times, "Host");
	
	// Run kernel on device to make sure everything is loaded
	Bfield_Map_t *checksums_d;
	cudaMalloc(&checksums_d, sizeof(Bfield_Map_t));
	TestAccessSpeedOnDevice<<<1,1>>>(bfield_d, 100, checksums_d);
	
	// Test access speed on device
	cout<<"Testing "<<Naccesses/1000000<<"M accesses on device"<<endl;
	cudaDeviceSynchronize();
	prof_time start_time_device;
	TestAccessSpeedOnDevice<<<1,1>>>(bfield_d, Naccesses, checksums_d);
	cudaDeviceSynchronize();
	start_time_device.TimeDiffNow(prof_times, "Device");
	
	// Copy results from device checksum to host
	Bfield_Map_t checksums_d_h;
	cudaMemcpy(&checksums_d_h, checksums_d, sizeof(Bfield_Map_t), cudaMemcpyDeviceToHost);
	
	// Print checksum info to screen
	cout<<"Checksums (first is host, second device): ----"<<endl;
	cout<<"r:  "<<checksums_h->r<<endl;
	cout<<"    "<<checksums_d_h.r<<endl;
	cout<<endl;
	cout<<"z:  "<<checksums_h->z<<endl;
	cout<<"    "<<checksums_d_h.z<<endl;
	cout<<endl;
	cout<<"Br  "<<checksums_h->Br<<endl;
	cout<<"    "<<checksums_d_h.Br<<endl;
	cout<<endl;
	cout<<"Bz: "<<checksums_h->Bz<<endl;
	cout<<"    "<<checksums_d_h.Bz<<endl;
	cout<<endl;

	// Print timing report to screen
	cout<<endl;
	cout<<"Profiling info ------"<<endl;
	map<string, prof_time::time_diffs>::iterator iter = prof_times.begin();
	for(; iter!=prof_times.end(); iter++){
		cout<< " "<<iter->first<<endl
		    <<"   real="<< iter->second.real*1000.0<<"ms"<<endl
		    <<"   prof="<< iter->second.prof*1000.0<<"ms"<<endl
		    <<"   virt="<< iter->second.virt*1000.0<<"ms"<<endl
		    << endl;
	}
	cout<<endl;
 	
	// Free memory on device and host
	cout<<"Freeing memory"<<endl;
	cudaFree(bfield_d);
	cudaFree(checksums_d);
	free(bfield_h);
	free(checksums_h);

	return 0;
}

//--------------------------
// TestAccessSpeedOnHost
//--------------------------
void TestAccessSpeedOnHost(Bfield_Map_t* bfield, int Naccesses, Bfield_Map_t *checksums)
{
	//printf("Hello from host!\n");
	checksums->r = 0.0;
	checksums->z = 0.0;
	checksums->Br = 0.0;
	checksums->Bz = 0.0;
	checksums->dBrdr = 0.0;
	checksums->dBzdr = 0.0;
	checksums->dBrdz = 0.0;
	checksums->dBzdz = 0.0;

	float delta_r = (BR_MAX-BR_MIN)/(float)(BR_N-1);
	float delta_z = (BZ_MAX-BZ_MIN)/(float)(BZ_N-1);
	for(int i=0; i<Naccesses; i++){
		float r = BR_MIN + (float)i*(BR_MAX-BR_MIN)/(float)Naccesses;
		float z = BZ_MIN + (float)i*(BZ_MAX-BZ_MIN)/(float)Naccesses;

		int ir = (int)((r - BR_MIN)/delta_r);
		int iz = (int)((z - BZ_MIN)/delta_z);
		Bfield_Map_t *B = &bfield[iz*BR_N + ir];

		Bfield_Map_t my_B = *B;
		checksums->r += r;
		checksums->z += z;
		checksums->Br += my_B.Br;
		checksums->Bz += my_B.Bz;
		checksums->dBrdr += my_B.dBrdr;
		checksums->dBzdr += my_B.dBzdr;
		checksums->dBrdz += my_B.dBrdz;
		checksums->dBzdz += my_B.dBzdz;
	}
}

//--------------------------
// TestAccessSpeedOnDevice
//--------------------------
__global__
void TestAccessSpeedOnDevice(Bfield_Map_t* bfield, int Naccesses, Bfield_Map_t *checksums)
{
	//printf("Hello from device!\n");
	checksums->r = 0.0;
	checksums->z = 0.0;
	checksums->Br = 0.0;
	checksums->Bz = 0.0;
	checksums->dBrdr = 0.0;
	checksums->dBzdr = 0.0;
	checksums->dBrdz = 0.0;
	checksums->dBzdz = 0.0;

	float delta_r = (BR_MAX-BR_MIN)/(float)(BR_N-1);
	float delta_z = (BZ_MAX-BZ_MIN)/(float)(BZ_N-1);
	for(int i=0; i<Naccesses; i++){
		float r = BR_MIN + (float)i*(BR_MAX-BR_MIN)/(float)Naccesses;
		float z = BZ_MIN + (float)i*(BZ_MAX-BZ_MIN)/(float)Naccesses;

		int ir = (int)((r - BR_MIN)/delta_r);
		int iz = (int)((z - BZ_MIN)/delta_z);
		Bfield_Map_t *B = &bfield[iz*BR_N + ir];

		Bfield_Map_t my_B = *B;
		checksums->r += r;
		checksums->z += z;
		checksums->Br += my_B.Br;
		checksums->Bz += my_B.Bz;
		checksums->dBrdr += my_B.dBrdr;
		checksums->dBzdr += my_B.dBzdr;
		checksums->dBrdz += my_B.dBrdz;
		checksums->dBzdz += my_B.dBzdz;
	}
}