#include <iostream>
#include <vector>
#include "DKSBase.h"

using namespace std;

void initData(vector< vector<double> > &v, int length) {

  for (unsigned int i = 0; i < v.size(); i++) {
    for (int j = 0; j < length; j++) {
      v[i].push_back(j);
    }
  }

}


void printData(vector< vector<double> > &v) {
  for (unsigned int i = 0; i < v.size(); i++) {
    for (unsigned int j = 0; j < v[i].size(); j++) {
      cout << v[i][j] << "\t";
    }
    cout << endl;
  }
}

void initData(double *data, int sensors, int length) {

  for (int i = 0; i < sensors; i++) {
    for (int j = 0; j < length; j++) {
      data[i*length + j] = j;
    }
  }

}


void printData(double *data, int sensors, int length) {
  for (int i = 0; i < sensors; i++) {
    for (int j = 0; j < length; j++) {
      cout << data[i*length + j] << "\t";
    }
    cout << endl;
  }
}

void initPar(double *par, int npar) {

  for (int i = 0; i < npar; i++)
    par[i] = (double)i / npar;

}

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

void calcChisq(vector< vector<double> > fData, double * par, double fTimeResolution, double fRebin)
{

  double chisq = 0.0;
  double theo, data;
  const double tau=2.197019;
  const double dt0 = fTimeResolution*0.5*(fRebin-1);
  double time;
  double w = par[0]*0.08516155035269027; 

  unsigned int i, j;

  for (i=0; i<fData.size(); i++) {
    for (j=0; j<fData[0].size(); j++) {
      data = fData[i][j];
      time = dt0+fTimeResolution*fRebin*j;

      theo = par[2 + i*4] * exp(-time/tau)*(1.0 + par[3 + i*4]*exp(-0.5 * pow(par[1]*time,2.0))*cos(w*time+par[4+i*4]*1.74532925199432955e-2))+par[5+i*4]; 
      if (data != 0.0) {
	chisq += (theo-data)*(theo-data)/data;
	cout << (theo-data)*(theo-data)/data << "\t";
      } else {
	chisq += theo*theo;
	cout << theo*theo << "\t";
      }
    }
    cout << endl;
  }

  cout << "Chisq: " << chisq << endl;

}


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

  bool useCuda = true;
  if (argc == 2 && atoi(argv[1]) == 1)
    useCuda = false;

  int ierr;
  int sensors = 5;
  int length = 10;
  int npar = 4 * sensors + 2;
  int ndata = sensors * length;
  
  double result;

  double fTimeResolution = 0.05;
  double fRebin = 5;

  double *par = new double[npar];
  initPar(par, npar);
  
  vector< vector< double > > fData;
  fData.resize(sensors);
  initData(fData, length);
  printData(fData);
  printDiv(75);

  DKSBase dksbase;
  if (useCuda)
    dksbase.setAPI("Cuda", 4);
  else
    dksbase.setAPI("OpenCL", 6);
  dksbase.setDevice("-gpu", 4);
  dksbase.initDevice();
  dksbase.setupFFT(0, NULL);


  void *mem_data, *mem_par, *mem_chisq;
  cout << "Allocate memory" << endl;
  mem_par = dksbase.allocateMemory<double>(npar, ierr);
  mem_data = dksbase.allocateMemory<double>(fData.size() * fData[0].size(), ierr);
  mem_chisq = dksbase.allocateMemory<double>(fData.size() * fData[0].size(), ierr);
  
  
  cout << "Write data" << endl;
  dksbase.writeData<double>(mem_par, par, npar);
  for (int i = 0; i < sensors; i++)
    dksbase.writeData<double>(mem_data, &fData[i][0], length, i*length);
  
  
  
  cout << "Call PHistoTFFcn" << endl;
  dksbase.callPHistoTFFcn(mem_data, mem_par, mem_chisq, 
			  fTimeResolution, fRebin, 
			  sensors, length, npar, result);
  cout << "Result: " << result << endl;
  
  
  double *out_data = new double[ndata];
  dksbase.readData<double>(mem_chisq, out_data, ndata);
  printDiv(75);
  printData(out_data, sensors, length);
  printDiv(75);
  
  calcChisq(fData, par, fTimeResolution, fRebin);
  printDiv(75);
  
  cout << "Free memory" << endl;
  dksbase.freeMemory<double>(mem_par, npar);
  dksbase.freeMemory<double>(mem_data, ndata);
  dksbase.freeMemory<double>(mem_chisq, ndata);
  

  return 0;

}