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

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

using namespace std;

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

void initData(complex<double> *data, int dimsize[3], int dim);
void initData(double *data, int dimsize[3], int dim);

bool readParams(int argc, char *argv[], int &N1, int &N2, int &N3, int &dim, 
		char *api_name, char *device_name);

void printHelp();

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

  int ierr;
  int N1 = 8;
  int N2 = 8;
  int N3 = 8;
  int dim = 3;
  char *api_name = new char[10];
  char *device_name = new char[10];

  if ( readParams(argc, argv, N1, N2, N3, dim, api_name, device_name) )
    return 0;

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

  int dimsize[3] = {N1, N2, N3};
  int sizereal = dimsize[0] * dimsize[1] * dimsize[2];
  int sizecomp = (dimsize[0]/2+1) * dimsize[1] *dimsize[2];

  double *rdata = new double[sizereal];
  double *ordata = new double[sizereal];
  complex<double> *cdata = new complex<double>[sizereal];
  complex<double> *codata = new complex<double>[sizereal];

  initData(rdata, dimsize, 3);
  initData(cdata, dimsize, 3);

  /* init DKSBase */
  cout << "Init device and set function" << endl;
  DKSFFT base;
  base.setAPI(api_name, strlen(api_name));
  base.setDevice(device_name, strlen(device_name));
  cout << "init device" << endl;
  base.initDevice();
  cout << "setup fft" << endl;
  base.setupFFT(dim, dimsize);

  //Test RC FFT -> CR FFT
  void *real_ptr, *comp_ptr, *res_ptr;
  cout << "allocate memory" << endl;
  real_ptr = base.allocateMemory<double>(sizereal, ierr);
  res_ptr = base.allocateMemory<double>(sizereal, ierr);
  comp_ptr = base.allocateMemory< complex<double> >(sizecomp, ierr);

  cout << "write data" << endl;
  base.writeData<double>(real_ptr, rdata, sizereal);

  cout << "perform fft" << endl;
  base.callR2CFFT(real_ptr, comp_ptr, dim, dimsize);
  base.callC2RFFT(res_ptr, comp_ptr, dim, dimsize);
  base.callNormalizeC2RFFT(res_ptr, dim, dimsize);

  cout << "read data" << endl;
  base.readData<double>(res_ptr, ordata, sizereal);

  compareData(rdata, ordata, N1, 3);

  base.freeMemory<double>(real_ptr, sizereal);
  base.freeMemory<double>(res_ptr, sizereal);
  base.freeMemory< complex<double> >(comp_ptr, sizecomp);

  //Test CC FFT
  void *mem_ptr;
  mem_ptr = base.allocateMemory< complex<double> >(sizereal, ierr);
  base.writeData< complex<double> >(mem_ptr, cdata, sizereal);
  base.callFFT(mem_ptr, 3, dimsize);
  base.callIFFT(mem_ptr, 3, dimsize);
  base.callNormalizeFFT(mem_ptr, 3, dimsize);
  base.readData< complex<double> >(mem_ptr, codata, sizereal);

  compareData(cdata, codata, N1, 3);

  base.freeMemory< complex<double> > (mem_ptr, sizereal);

  delete[] rdata;
  delete[] ordata;
  delete[] cdata;
  delete[] codata;

}

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

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 << " RC <--> CR diff: " << sum << endl;
}

void initData(complex<double> *data, int dimsize[3], int dim) {
  if (dim == 3) {
    for (int i = 0; i < dimsize[2]; i++)
      for (int j = 0; j < dimsize[1]; j++) 
	for (int k = 0; k < dimsize[0]; k++) 
	  data[i*dimsize[1]*dimsize[0] + j*dimsize[0] + k] = complex<double>(sin(k), 0.0);
  } else if (dim == 2) {
    for (int j = 0; j < dimsize[1]; j++) {
      for (int k = 0; k < dimsize[0]; k++) {
	data[j*dimsize[0] + k] = complex<double>(sin(k), 0.0);
      }
    }
  } else {
    for (int k = 0; k < dimsize[0]; k++) 
      data[k] = complex<double>(sin(k), 0.0);
  }
}

void initData(double *data, int dimsize[3], int dim) {
  if (dim == 3) {
    for (int i = 0; i < dimsize[2]; i++)
      for (int j = 0; j < dimsize[1]; j++) 
	for (int k = 0; k < dimsize[0]; k++) 
	  data[i*dimsize[1]*dimsize[0] + j*dimsize[0] + k] = sin(k);
  } else if (dim == 2) {
    for (int j = 0; j < dimsize[1]; j++) {
      for (int k = 0; k < dimsize[0]; k++) {
	data[j*dimsize[0] + k] = sin(k);
      }
    }
  } else {
    for (int k = 0; k < dimsize[0]; k++) 
      data[k] = sin(k);
  }
}

bool readParams(int argc, char *argv[], int &N1, int &N2, int &N3, int &dim,
		char *api_name, char *device_name) 
{

  for (int i = 1; i < argc; i++) {

    if ( argv[i] == std::string("-dim")) {
      dim = atoi(argv[i + 1]);
      i++;
    }

    if ( argv[i] == std::string("-grid") ) {
      N1 = atoi(argv[i + 1]);
      N2 = atoi(argv[i + 2]);
      N3 = atoi(argv[i + 3]);
      i += 3;
    }

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

  return false;
}