#include <iostream>
#include <cstdlib>
#include <complex>

#include "Utility/TimeStamp.h"
#include "DKSBase.h"

using namespace std;

void compareData(complex<double>* &data1, complex<double>* &data2, int N, int dim);


int main(int argc, char *argv[]) {

	int N = 4;
	char *api_name = new char[10];
	char *device_name = new char[10];
	if (argc == 2) {
		strcpy(api_name, argv[1]);
		strcpy(device_name, "-gpu");
	} else if (argc > 2) {
		strcpy(api_name, argv[1]);
		strcpy(device_name, argv[2]);
		N = atoi(argv[3]);
	} else {
		strcpy(api_name, "OpenCL");
		strcpy(device_name, "-gpu");
	}
	int dimsize[3] = {N, N, N};


	cout << "Use api: " << api_name << endl;

	cout << "Begin DKS Base tests, N = " <<  N << endl;		

	complex<double> *cdata = new complex<double>[N*N*N];
	complex<double> *cfft = new complex<double>[N*N*N];
	complex<double> *cifft = new complex<double>[N*N*N];
	
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < N; j++) {
			for (int k = 0; k < N; k++) {
				cdata[i*N*N + j*N + k] = complex<double>((double)i / N, 0);
				cfft[i*N*N + j*N + k] = complex<double>(0, 0);
				cifft[i*N*N + j*N + k] = complex<double>(0, 0);
			}
		}
	}
	
	timestamp_t t0, t1;
	
	/* init DKSBase */
	cout << "Init device and set function" << endl;
	DKSBase base;
	base.setAPI(api_name, strlen(api_name));
	base.setDevice(device_name, strlen(api_name));
	base.initDevice();
	
	void *mem_ptr;
	int ierr;
	
	/* run stest funct to init device */
	mem_ptr = base.allocateMemory< complex<double> >(N*N*N, ierr);
	ierr = base.writeData< complex<double> >(mem_ptr, cdata, N*N*N);
	base.callFFT(mem_ptr, 3, dimsize);
	base.callIFFT(mem_ptr, 3, dimsize);
	base.callNormalizeFFT(mem_ptr, 3, dimsize);
	base.readData< complex<double> >(mem_ptr, cifft, N*N*N);
	base.freeMemory< complex<double> >(mem_ptr, N*N*N);
	/* end test */
	
	int steps = 10;
	base.oclClearEvents();
	t0 = get_timestamp();
	for (int i = 0; i < steps; i++) {
	
		/* allocate memory on device */
		mem_ptr = base.allocateMemory< complex<double> >(N*N*N, ierr);
		
		/* write data to device */	
		ierr = base.writeData< complex<double> >(mem_ptr, cdata, N*N*N);

		/* execute fft */
		base.callFFT(mem_ptr, 3, dimsize);
	
		/* execute ifft */	
		base.callIFFT(mem_ptr, 3, dimsize);
	
		/* execute normalize */
		base.callNormalizeFFT(mem_ptr, 3, dimsize);
	
		/* read data from device */
		base.readData< complex<double> >(mem_ptr, cifft, N*N*N);
		
		/* free device memory */
		base.freeMemory< complex<double> >(mem_ptr, N);

		//compareData(cdata, cifft, N, 3);
	}
	t1 = get_timestamp();
	
	cout << "=========================" << endl;
	//base.oclEventInfo();
	cout << "Average total: " << get_secs(t0, t1) / steps << endl;
	cout << "=========================" << endl;
	
	
	
		
	return 0;
}

void compareData(complex<double>* &data1, complex<double>* &data2, int N, int dim) {
    int ni, nj, nk, id;
    ni = (dim > 2) ? N : 1;
    nj = (dim > 1) ? N : 1;
    nk = N;
    double sum = 0;
    for (int i = 0; i < ni; i++) {
		for (int j = 0; j < nj; j++) {
	    	for (int k = 0; k < nk; k++) {
			id = i*ni*ni + j*nj + k;
			sum += fabs(data1[id].real() - data2[id].real());
			sum += fabs(data1[id].imag() - data2[id].imag());
		    }
		}
    }
    cout << "Size " << N << " CC <--> CC diff: " << sum << endl;
}