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

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

using namespace std;

void printData(complex<double>* &data, int N, int dim, bool normalize = false);
void printData3DN4(complex<double>* &data, int N, int dim);

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

/* usage - ./testFFT3D */
int main(int argc, char *argv[]) {

  int N = 16;
  char *api_name = new char[10];
  char *device_name = new char[10];

  for (int i = 1; i < argc; i++) {
    if (argv[i] == string("-cuda")) {
      strcpy(api_name, "Cuda");
      strcpy(device_name, "-gpu");
    } 

    if (argv[i] == string("-opencl")) {
      strcpy(api_name, "OpenCL");
      strcpy(device_name, "-gpu");
    } 

    if (argv[i] == string("-mic")) {
      strcpy(api_name, "OpenMP");
      strcpy(device_name, "-mic");
    } 

    if (argv[i] == string("-cpu")) {
      strcpy(api_name, "OpenCL");
      strcpy(device_name, "-cpu");
    }

    if (argv[i] == string("-N"))
      N = atoi(argv[i+1]);
  }

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

  int dimsize[3] = {N, N, N};

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

  int dim = 3;
  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)k / N, 0);
	cfft[i*N*N + j*N + k] = complex<double>(0, 0);
	cifft[i*N*N + j*N + k] = 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(device_name));
  base.initDevice();
  base.setupFFT(3, dimsize);
	
  void *mem_ptr;
  int ierr;

  /* 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);
  if (N < 5)
    printData3DN4(cdata, N, 3);


  /* execute fft */
  base.callFFT(mem_ptr, 3, dimsize);
  if (N < 5) {
    base.readData< complex<double> > (mem_ptr, cfft, N*N*N);
    printData3DN4(cfft, N, 3);
  }
	
  /* 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);
  if (N < 5)
    printData3DN4(cifft, N, 3);	

  /* free device memory */
  base.freeMemory< complex<double> >(mem_ptr, N*N*N);
	
  /* compare results */
  compareData(cdata, cifft, N, dim);
		
  return 0;
}

void printData(complex<double>* &data, int N, int dim, bool normalize) {
  int ni, nj, nk;
  ni = (dim > 2) ? N : 1;
  nj = (dim > 1) ? N : 1;
  nk = N;
    
  for (int i = 0; i < ni; i++) {
    for (int j = 0; j < nj; j++) {
      for (int k = 0; k < nk; k++) {
	if (!normalize) {
	  cout << data[i*ni*ni + j*nj + k].real() << " ";
	  cout << data[i*ni*ni + j*nj + k].imag() << "\t";
	} else
	  cout << data[i*ni*ni + j*nj + k].real() / N << "\t";
      }
      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 << ") ";
      }
    }
    cout << endl;
  }
  cout << endl;
    
}

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;
}