#include <iostream>

#include <vector>

#include "DKSBase.h"
#include "cuda_runtime.h"

#include <mpi.h>

using namespace std;

typedef struct {
  int label;
  unsigned localID;
  double Rincol[3];
  double Pincol[3];
  long IDincol;
  int Binincol;
  double DTincol;
  double Qincol;
  long LastSecincol;
  double Bfincol[3];
  double Efincol[3];
} PART;

PART initPart(int d) {

  PART p;
  p.label = d;
  p.localID = d;
  for (int i = 0; i < 3; i++) {
    p.Rincol[i] = 0.5;// / (d+1);
    p.Pincol[i] = 0.5;// / (d+1);
    p.Bfincol[i] = 1.0 / (d+1);
    p.Efincol[i] = 1.0 / (d+1);
  }
  p.IDincol = d;
  p.Binincol = d;
  p.DTincol = d;
  p.Qincol = d;
  p.LastSecincol = d;

  return p;

}

void printPart(PART p) {

  cout << "label: " << p.label << ", ";
  //cout << "localID: " << p.localID << ", ";
  cout << "Rincol: " << p.Rincol[0] << ", " << p.Rincol[1] << ", " << p.Rincol[2] << ", ";
  cout << "Pincol: " << p.Pincol[0] << ", " << p.Pincol[1] << ", " << p.Pincol[2] << ", ";
  //cout << "IDincol: " << p.IDincol << ", Binincol: " << p.Binincol << ", ";
  //cout << "DTincol: " << p.DTincol << ", Qincol: " << p.Qincol << ", LastSecincol: " << p.LastSecincol << ", ";
  //cout << "Bfincol: " << p.Bfincol[0] << ", " << p.Bfincol[1] << ", " << p.Bfincol[2] << ", ";
  //cout << "Efincol: " << p.Efincol[0] << ", " << p.Efincol[1] << ", " << p.Efincol[2] << endl;
  cout << endl;


}

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

  int ierr;
  int rank, nprocs;
		
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

  int numpart = 500501;

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

  PART tmp;
  vector<PART> p;
  vector<PART> p_out;
  p_out.resize(numpart);

  for (int i = 0; i < numpart; i++) {
    tmp = initPart(i + 1);
    p.push_back(tmp);
  }

  if (numpart <= 20) {
    for (int i = 0; i < 10; i++)
      printPart(p[i]);
    cout << endl;
  }

  double params[19];
  for (int i = 0; i < 19; i++)
    params[i] = 0.05;
  params[0] = 0;
  params[1] = 1;

  void *mem_ptr, *par_ptr;
  
  par_ptr = base.allocateMemory<double>(19, ierr);
  base.writeData<double>(par_ptr, params, 19);

  mem_ptr = base.allocateMemory<PART>(numpart, ierr);
  base.writeData<PART>(mem_ptr, &p[0], numpart);

  int addback, dead;
  for (int i = 0; i < 100; i++)
    base.callCollimatorPhysics(mem_ptr, par_ptr, numpart, 19, addback, dead);
  cout << "Add back: " << addback << ", dead: " << dead << endl;

  base.readData<PART>(mem_ptr, &p_out[0], numpart);
  base.freeMemory<PART>(mem_ptr, ierr);
  base.freeMemory<double>(par_ptr, ierr);

  if (numpart <= 20) {
    for (int i = 0; i < numpart; i++)
      printPart(p_out[i]);
  }

  MPI_Finalize();
  return 0;

}