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

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

using namespace std;

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

	int n = 2;
	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 == 3) {
		strcpy(api_name, argv[1]);
		strcpy(device_name, argv[2]);
	} else if (argc == 4) {
		strcpy(api_name, argv[1]);
		strcpy(device_name, argv[2]);
		n = atoi(argv[3]);
	} else {
		strcpy(api_name, "OpenCL");
		strcpy(device_name, "-gpu");
	}

	int N = pow(2,n);
	cout << "Use api: " << api_name << endl;

	cout << "Begin DKS Base tests" << endl;
		
	cout << "FFT size: " << N << endl;
	
	int dimsize[3] = {N, N, N};
	
	complex<double> *cdata = new complex<double>[N];
	complex<double> *cfft = new complex<double>[N];
	complex<double> *cfft2 = new complex<double>[N];
	complex<double> *cfftsrc = new complex<double>[N];
	for (int i = 0; i < N; i++) {
		cdata[i] = complex<double>((double)i / N, 0);
		cfft[i] = complex<double>(0, 0);
		cfft2[i] = complex<double>(0, 0);
		cfftsrc[i] = complex<double>(0, 0);
	}
	
	/* 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();
	
	
	timestamp_t t0, t1;
	
	/* radix-2 in place fft */
	void *mem_ptr;
	int ierr;
	
	for (int i = 0; i < 5; i++) {
		t0 = get_timestamp();
		mem_ptr = base.allocateMemory< complex<double> >(N, ierr);
		base.writeData< complex<double> >(mem_ptr, cdata, N);
		base.callFFT(mem_ptr, 1, dimsize);
		base.readData< complex<double> >(mem_ptr, cfft, N);
		base.freeMemory< complex<double> >(mem_ptr, N);
		t1 = get_timestamp();
		cout << "in-place FFT time: " << get_secs(t0, t1) << endl;
	}
	
	cout << endl;
	
	/* stockham radix-2 out-of-place fft */
	void *src_ptr;
	
	for (int i = 0; i < 5; i++) {
		t0 = get_timestamp();
		src_ptr = base.allocateMemory< complex<double> >(N, ierr);
		base.writeData< complex<double> >(src_ptr, cdata, N);
		base.callFFTStockham(src_ptr, 1, dimsize);
		base.readData< complex<double> >(src_ptr, cfft2, N);
		base.freeMemory< complex<double> >(src_ptr, N);
		t1 = get_timestamp();
		cout << "out-of-place FFT time: " << get_secs(t0, t1) << endl;
	}
	
	double diff = 0;
	for (int i = 0; i < N; i++) {
		diff += fabs(cfft[i].real() - cfft2[i].real());
		diff += fabs(cfft[i].imag() - cfft2[i].imag());
	}
	
	cout << endl << "Difference: " << diff << endl;
	
	if (diff > 0.00001) {
		for (int i = 0; i < 10; i++) {
			cout << cfft[i] << "\t" << cfft2[i] << endl;
		}
	}
		
	return 0;
}