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

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

using namespace std;

void printData(double* &data, int N1, int N2);
void printData(complex<double>* &data, int N1, int N2);
void printData3DN4(complex<double>* &data, int N, int dim);
void printData3DN4(double* &data, int N, int dim);


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



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

	int N1 = 4;
	int N2 = 4;

	if (argc == 3) {
		N1 = atoi(argv[1]);
		N2 = atoi(argv[2]);
	}

	int dimsize[3] = {N1, N2, 1};

	cout << "Begin RC 3D FFT tests, grid = " <<  N1 << "\t" << N2 << endl;		
	int sizereal = N1*N2;
	int sizecomp = N1*(N2/2+1);

	int dim = 3;
	double *cdata = new double[sizereal];
	complex<double> *cfft = new complex<double>[sizecomp];
	
	for (int i = 0; i < N2; i++) {
		for (int j = 0; j < N1; j++) {
			cdata[i*N1 + j] = (double)(j) / N1;
		}
	}
	
	/* init DKSBase */
	cout << "Init device and set function" << endl;
	DKSBase base;
	base.setAPI("Cuda", 4);
	base.setDevice("-gpu", 4);
	base.initDevice();
	
	void *real_ptr, *comp_ptr;
	int ierr;
	/* allocate memory on device */
	real_ptr = base.allocateMemory<double>(sizereal, ierr);
	comp_ptr = base.allocateMemory< complex<double> >(sizecomp, ierr);
	
	/* write data to device */	
	ierr = base.writeData<double>(real_ptr, cdata, sizereal);

	/* execute fft */
	base.callR2CFFT(real_ptr, comp_ptr, 2, dimsize);
	
	/* read data from device */
	base.readData< complex<double> >(comp_ptr, cfft, sizecomp);
	
	/* free device memory */
	base.freeMemory<double>(real_ptr, sizereal);
	base.freeMemory< complex<double> >(comp_ptr, sizecomp);
	
	cout << "FFT complete" << endl;
	
	
	/* print results */
	printData(cdata, N1, N2);
	printData(cfft, N1, N2);
	
	
		
	return 0;
}

void printData(double* &data, int N1, int N2) {
    
    for (int i = 0; i < N2; i++) {
		for (int j = 0; j < N1; j++) {
			cout << data[i*N1 + j] << " ";
		}
		cout << endl;
    }
	cout << endl;
}

void printData(complex<double>* &data, int N1, int N2) {
    
    complex<double> tmp(0.0, 0.0);
    for (int i = 0; i < N2/2+1; i++) {
		for (int j = 0; j < N1; j++) {
			tmp = data[i*N1 + j];
		    if (tmp.real() < 0.00001 && tmp.real() > -0.00001) tmp = complex<double>(0.0, tmp.imag());
		    if (tmp.imag() < 0.00001 && tmp.imag() > -0.00001) tmp = complex<double>(tmp.real(), 0.0);
		    	
		    cout << tmp << " ";
		}
		cout << endl;
    }
    cout << endl;
}

void printData3DN4(complex<double>* &data, int N, int dim) {
    
	for (int j = 0; j < N; j++) {
		for (int i = 0; i < N; i++) {
			for (int k = 0; k < N; k++) {
				
				double d = data[i*N*N + j*N + k].real();
				double a = data[i*N*N + j*N + k].imag();
				
				if (d < 10e-5 && d > -10e-5)
					d = 0;
				if (a < 10e-5 && a > -10e-5)
					a = 0;
					
				cout << d << "; " << a << "\t";
    		}
    	}
	    cout << endl;
	}
	cout << endl;
    
}

void printData3DN4(double* &data, int N, int dim) {
    
	for (int j = 0; j < N; j++) {
		for (int i = 0; i < N; i++) {
			for (int k = 0; k < N; k++) {
				double d = data[i*N*N + j*N + k];
				if (d > 10e-5 || d < -10e-5)
				    cout << d << "\t";
				else 
					cout << 0 << "\t";
    		}
    	}
	    cout << endl;
	}
	cout << endl;
    
}

void compareData(double* &data1, 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] - data2[id]);
		    }
		}
    }
    cout << "Size " << N << " CC <--> CC diff: " << sum << endl;
}