#include <iostream>
#include <mpi.h>
#include <string.h>
#include <complex>

#include "DKSBase.h"
#include "nvToolsExt.h"
#include "cuda_profiler_api.h"
#include "cuda_runtime.h"

using namespace std;


void printData3D(double* data, int N, int NI, const char *message = "") {
	if (strcmp(message, "") != 0)
		cout << message;
   
  for (int i = 0; i < NI; i++) {
	  for (int j = 0; j < N; j++) {
	    for (int k = 0; k < N; k++) {
		    cout << data[i*N*N + j*N + k] << "\t";
	    }
		  cout << endl;
		}
	  cout << endl;
  }
    
}

void initData(double *data, int N) {

  for (int i = 0; i < N/4 + 1; i++) {
    for (int j = 0; j < N/2 + 1; j++) {
      for (int k = 0; k < N/2 + 1; k++) {
        data[i*N*N + j*N + k] = k+1;
      }
    }
  }
}

void initData2(double *data, int N) {
  for (int i = 0; i < N; i++)
    data[i] = i;
}

void initComplex( complex<double> *d, int N) {

  for (int i = 0; i < N; i++) {
    d[i] = complex<double>(2, 0);
  }

}

void printComplex(complex<double> *d, int N) {
  
  for (int i = 0; i < N; i++)
    cout << d[i] << "\t";
  cout << endl;

}

void initMirror(double *data, int n1, int n2, int n3) {
  int d = 1;
  for (int i = 0; i < n3; i++) {
    for (int j = 0; j < n2; j++) {
      for (int k = 0; k < n1; k++) {
	if (i < n3/2 + 1 && j < n2/2 + 1 && k < n1/2 + 1)
	  data[i * n2 * n1 + j * n1 + k] = d++;
	else
	  data[i * n2 * n1 + j * n1 + k] = 0;
      }
    }
  }
}

void printDiv(int c) {
  for (int i = 0; i < c; i++)
    cout << "-";
  cout << endl;

}

void printMirror(double *data, int n1, int n2, int n3) {
  
  printDiv(75);
  for (int i = 0; i < n3; i++) {
    for (int j = 0; j < n2; j++) {
      for (int k = 0; k < n1; k++) {
	cout << data[i * n2 * n1 + j * n1 + k] << "\t";
      }
      cout << endl;
    }
    cout << endl;
  }
  cout << endl;
}

double sumData(double *data, int datasize) {

  double sum = 0;
  for (int i = 0; i < datasize; i++)
    sum += data[i];

  return sum;
}



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

  int ierr;
 
  int N1 = 8;
  int N2 = 8;
  int N3 = 4;

  int n1 = N1 / 2; 
  int n2 = N2 / 2;
  int n3 = N3 / 2;

  int sizegreen = (n1 + 1) * (n2 + 1) * (n3 + 1);
  int sizerho = N1 * N2 * N3;

  double *data_green; //= new double[sizegreen];
  double *data_rho; //= new double[sizerho];

  double hr_m0 = +4.0264984513873269e-04;
  double hr_m1 = +4.3305596731911289e-04;
  double hr_m2 = +8.3154085085560838e-04;

  DKSBase base = DKSBase();
  base.setAPI("Cuda", 4);
  base.setDevice("-gpu", 4);
  base.initDevice();

  
  int stream1, stream2;
  base.createStream(stream1);
  base.createStream(stream2);
  cout << "ID stream1: " << stream1 << endl;
  cout << "ID stream2: " << stream2 << endl;
  
  void *mem_green1, *mem_green2, *mem_rho1, *mem_rho2;

  mem_green1 = base.allocateMemory<double>(sizegreen, ierr);
  mem_green2 = base.allocateMemory<double>(sizegreen, ierr);
  mem_rho1 = base.allocateMemory<double>(sizerho, ierr);
  mem_rho2 = base.allocateMemory<double>(sizerho, ierr);

  printDiv(50);

  data_green = new double[sizegreen];
  data_rho = new double[sizerho];
  
  base.callGreensIntegral(mem_green1, n1+1, n2+1, n3+1, n1+1, n2+1, 
			  hr_m0, hr_m1, hr_m2, stream1);
  base.readData<double>(mem_green1, data_green, sizegreen);
  cout << "Sum green: " << sumData(data_green, sizegreen) << endl;
  cout << scientific << setprecision(16);
  for (int p = 0; p < 7; p++)
    cout << data_green[p] << "\t";
  cout << endl;
  //printMirror(data_green, n1 + 1, n2 + 1, n3 + 1);

  base.callGreensIntegration(mem_rho1, mem_green1, n1 + 1, n2 + 1, n3 + 1, -1);
  base.readData<double>(mem_rho1, data_rho, sizerho);
  cout << "Sum integral: " << sumData(data_rho, sizerho) << endl;
  //printMirror(data_rho, N1, N2, N3);

  base.callMirrorRhoField(mem_rho1, n1, n2, n3, -1);
  base.readData<double>(mem_rho1, data_rho, sizerho);
  cout << "Sum mirror: " << sumData(data_rho, sizerho) << endl;
  //printMirror(data_rho, N1, N2, N3);

  printDiv(50);

  /*
  base.callGreensIntegral(mem_green2, n1+1, n2+1, n3+1, n1+1, n2+1, 
			  1, 1, 1, -2);
  base.readData<double>(mem_green2, data_green, sizegreen);
  cout << "Sum green: " << sumData(data_green, sizegreen) << endl;
  //printMirror(data_green, n1 + 1, n2 + 1, n3 + 1);

  base.callGreensIntegration(mem_rho2, mem_green2, n1 + 1, n2 + 1, n3 + 1, -2);
  base.readData<double>(mem_rho2, data_rho, sizerho);
  cout << "Sum integral: " << sumData(data_rho, sizerho) << endl;
  //printMirror(data_rho, N1, N2, N3);

  base.callMirrorRhoField(mem_rho2, n1, n2, n3, -2);
  base.readData<double>(mem_rho2, data_rho, sizerho);
  cout << "Sum mirror: " << sumData(data_rho, sizerho) << endl;
  //printMirror(data_rho, N1, N2, N3);
  */
  printDiv(50);

  base.freeMemory<double>(mem_green1, sizegreen);
  base.freeMemory<double>(mem_green2, sizegreen);
  base.freeMemory<double>(mem_rho1, sizerho);
  base.freeMemory<double>(mem_rho2, sizerho);
  
  delete [] data_green;
  delete [] data_rho;

  //test complex multiplication
  int compsize = 300;
  complex<double> *data1 = new complex<double>[compsize];
  complex<double> *data2 = new complex<double>[compsize];
  for (int i = 0; i < compsize; i++) {
    data1[i] = complex<double>(i+1, i+2);
    data2[i] = complex<double>(i+3, i+4);
  }

  for (int i = 0; i < 3; i++)
    cout << data1[i] << "\t";
  cout << endl;
  for (int i = 0; i < 3; i++)
    cout << data2[i] << "\t";
  cout << endl;

  void *ptr1, *ptr2;
  ptr1 = base.allocateMemory< complex<double> >(compsize, ierr);
  ptr2 = base.allocateMemory< complex<double> >(compsize, ierr);

  base.writeData< complex<double> >(ptr1, data1, compsize);
  base.writeData< complex<double> >(ptr2, data2, compsize);

  base.callMultiplyComplexFields(ptr1, ptr2, compsize);
  
  base.readData< complex<double> >(ptr1, data1, compsize);

  for (int i = 0; i < 3; i++)
    cout << data1[i] << "\t";
  cout << endl;

  base.freeMemory< complex<double> >(ptr1, compsize);
  base.freeMemory< complex<double> >(ptr2, compsize);
  		
  return 0;
}