#include <iostream>
#include <vector>
#include <string>

#include "DKSOPAL.h"

#include <vector_types.h>
#include "cuda_runtime.h"

void initData(double3 *data, int N) {
  for (int i = 0; i < N; i++) {
    data[i].x = (double)rand() / RAND_MAX;
    data[i].y = (double)rand() / RAND_MAX;
    data[i].z = (double)rand() / RAND_MAX;
  }
}

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

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

  int loop = 10;
  int numpart = 1e5;
  char *api_name = new char[10];
  char *device_name = new char[10];
  strcpy(api_name, "Cuda");
  strcpy(device_name, "-gpu");

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

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

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

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

  }

  std::cout << "=========================BEGIN TEST=========================" << std::endl;
  std::cout << "Use api: " << api_name << "\t" << device_name << std::endl;
  std::cout << "Number of particles: " << numpart << std::endl;
  std::cout << "Number of loops: " << loop << std::endl;
  std::cout << "------------------------------------------------------------" << std::endl;
  
  int ierr;
  DKSOPAL dksbase;
  dksbase.setAPI(api_name, strlen(api_name));
  dksbase.setDevice(device_name, strlen(api_name));
  ierr = dksbase.initDevice();
  if (ierr != DKS_SUCCESS)
    std::cout << "Error with init device!" << std::endl;

  double3 *R = new double3[numpart];
  double3 *P = new double3[numpart];
  double3 *Ef = new double3[numpart];
  double3 *Bf = new double3[numpart];
  double *dt = new double[numpart];
  
  initData(R, numpart);
  initData(P, numpart);
  initData(Ef, numpart);
  initData(Bf, numpart);
  initDt(dt, numpart);

  void *r_ptr, *p_ptr, *ef_ptr, *bf_ptr, *dt_ptr;
  
  r_ptr = dksbase.allocateMemory<double3>(numpart, ierr);
  p_ptr = dksbase.allocateMemory<double3>(numpart, ierr);
  ef_ptr = dksbase.allocateMemory<double3>(numpart, ierr);
  bf_ptr = dksbase.allocateMemory<double3>(numpart, ierr);
  dt_ptr = dksbase.allocateMemory<double>(numpart, ierr);
  

  dksbase.writeData<double3>(r_ptr, R, numpart);
  dksbase.writeData<double3>(p_ptr, P, numpart);
  dksbase.writeData<double3>(ef_ptr, Ef, numpart);
  dksbase.writeData<double3>(bf_ptr, Bf, numpart);
  dksbase.writeData<double>(dt_ptr, dt, numpart);

  for (int i = 0; i < loop; ++i)
    dksbase.callParallelTTrackerPush(r_ptr, p_ptr, dt_ptr, numpart, 1.0);


  std::cout << std::fixed << std::setprecision(4);
  for (int i = 0; i < 10; i++)
    std::cout << R[i].x << "\t" << R[i].y << "\t" << R[i].z << std::endl;

  std:: cout << "..." << std::endl;
  
  for (int i = numpart - 10; i < numpart; i++)
    std::cout << R[i].x << "\t" << R[i].y << "\t" << R[i].z << std::endl;

  std::cout << "============" << std::endl;

  dksbase.readData<double3>(r_ptr, R, numpart);

  std::cout << std::fixed << std::setprecision(4);
  for (int i = 0; i < 10; i++)
    std::cout << R[i].x << "\t" << R[i].y << "\t" << R[i].z << std::endl;

  std:: cout << "..." << std::endl;
  
  for (int i = numpart - 10; i < numpart; i++)
    std::cout << R[i].x << "\t" << R[i].y << "\t" << R[i].z << std::endl;

  dksbase.freeMemory<double3>(r_ptr, numpart);
  dksbase.freeMemory<double3>(p_ptr, numpart);
  dksbase.freeMemory<double3>(ef_ptr, numpart);
  dksbase.freeMemory<double3>(bf_ptr, numpart);
  dksbase.freeMemory<double>(dt_ptr, numpart);

  delete[] R;
  delete[] P;
  delete[] Ef;
  delete[] Bf;
  delete[] dt;


}