#include <iostream>
#include <sstream>
#include <chrono>
#include <iomanip>
#include <fstream>
#include <omp.h>
#include <TApplication.h>
#include <TCanvas.h>
#include <TROOT.h>

#include "Garfield/ComponentGrid.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/AvalancheMicroscopic.hh"
#include "Garfield/AvalancheMC.hh"
#include "Garfield/Random.hh"
#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/RandomEngineRoot.hh"

using namespace Garfield;

enum class BeamShape {
  Square,
  Circle
};

int main(int argc, char *argv[]) {
  //const int seed = 123456;
  //Garfield::RandomEngineRoot randomEngine(seed);
  //Garfield::Random::SetEngine(randomEngine);
  TApplication app("app", &argc, argv);
  auto start = std::chrono::high_resolution_clock::now();

  const bool debug = false;
  const bool storeEndpoints = false;

  // -- Gas parameters -- //
  const double temperature = 293.15; // [K]
  const double pressure = 760.; // [Torr]
  const double fN2 = 77.; // [%]
  const double fO2 = 21.; // [%]
  const double fAr = 1.; // [%]
  const double fH2O = 1.; // [%]

  // -- Chamber geometry -- //
  const double xChamberGap = 0.1; // [cm]
  const double yChamberLength = 2.; // [cm]
  const double zChamberLength = 2.; // [cm]
  const double highVoltage = 500.; // [V]

  // -- Pulse parameters -- //
  const int nPairs = 1e4; // [1]
  const double tStartPulse = 0.0; // [ns] start of pulse
  const double tEndPulse = 0.0; // [ns] end of pulse 10 μs of Protheus ONE
  const double beamCrossSection = 1.; // [cm^2]
  BeamShape shape = BeamShape::Square;  // Beam shape: Square or Circle

  // -- Drift parameters -- //
  const double tStartSimu = 0.;
  const double tEndSimu = 4600; // [ns] end of simulation (no more overlap region)
  const double dt = 50.; // [ns] density map update time
  const bool withIonRecombination = true;
  const bool withDensityMap = true;
  const bool withIonDiffusion = true;
  const double alpha = 1.72e-15; // [cm^3/ns] ion-ion recombination coefficient

  // -- Grid parameters -- //
  const double sigma = 0.0107566; // [cm]
  const double pad  = 3 * sigma;
  const int nx = 20, ny = 20, nz = 20;


  double densityTarget; // [cm^-3]
  for (int i = 1; i < argc; i++) {
    std::string arg = argv[i];

    if (arg == "--density" && i + 1 < argc) {
      densityTarget = std::stod(argv[++i]);
    } else if (arg == "--help") {
      std::cout << "Usage: ./AskedDensity [options]\n"
                << "Options:\n"
                << "  --density <val> in [cm^-3]\n";
      return 0;
    }
  }

  struct P {
    double x, y, z, t;
    size_t w;
  };

  MediumMagboltz gas;
  gas.SetComposition("O2", fO2, "Ar", fAr, "H2O", fH2O, "N2", fN2);
  gas.SetTemperature(temperature); // [K] 
  gas.SetPressure(pressure); // [Torr]
  gas.LoadIonMobility("IonMobility_N2+_N2.txt");
  gas.LoadNegativeIonMobility("NegIonMobility_O2-_air.txt");
  gas.Initialise(true);

  ComponentAnalyticField cmp;
  cmp.SetMedium(&gas);
  cmp.AddPlaneX(0, 0, "cathode");
  cmp.AddPlaneX(xChamberGap, highVoltage, "anode");

  const double side = std::sqrt(beamCrossSection);
  const double xMinGrid = 0., xMaxGrid = xChamberGap;
  const double yMinGrid = -(side/2. + pad);
  const double yMaxGrid = (side/2. + pad);
  const double zMinGrid = -(side/2. + pad);
  const double zMaxGrid = (side/2. + pad);

  ComponentGrid grid;
  grid.SetMesh(nx, ny, nz, xMinGrid, xMaxGrid, yMinGrid, yMaxGrid, zMinGrid,
                zMaxGrid);
  grid.SetMedium(&gas);
  grid.SetUniformElectricField(0., 0., 0.);
  //grid.SaveElectricField(&grid, "field_map.txt","XYZ");
  if (debug) {
    unsigned int nxx, nyy, nzz;
    double xmin, xmax, ymin, ymax, zmin, zmax;
    if (grid.GetMesh(nxx, nyy, nzz, xmin, xmax, ymin, ymax, zmin, zmax)) {
        std::cout << "Mesh dimensions:" << std::endl;
        std::cout << "  nx = " << nxx
                  << ", ny = " << nyy
                  << ", nz = " << nzz << std::endl;
        std::cout << "Bounding box:" << std::endl;
        std::cout << "  xmin = " << xmin << ", xmax = " << xmax << std::endl;
        std::cout << "  ymin = " << ymin << ", ymax = " << ymax << std::endl;
        std::cout << "  zmin = " << zmin << ", zmax = " << zmax << std::endl;
    } else {
        std::cerr << "Erreur : GetMesh a échoué." << std::endl;
    }
  }

  Sensor sensor;
  sensor.AddComponent(&cmp);
  sensor.AddComponent(&grid);
  sensor.SetArea(0, -yChamberLength/2, -zChamberLength/2, xChamberGap,
                 yChamberLength/2, zChamberLength/2);
  

  std::vector<P> electrons;
  std::vector<P> ions;
  std::vector<P> negions;
  int nAttaElectrons = 0;
  int nRecombIons = 0;
  int nRecombNegions = 0;

  const size_t multiplicity = densityTarget * beamCrossSection * xChamberGap
                            / nPairs; // multiplicity of pairs created
  for (int i = 0; i < nPairs; i++) {
    double t = tStartPulse + (static_cast<double>(i) / (nPairs - 1))
              * (tEndPulse - tStartPulse);
    double x0, y0, z0;
    switch (shape) {
      case BeamShape::Square: {
        x0 = Garfield::RndmUniform() * xChamberGap;
        y0 = (Garfield::RndmUniform() - 0.5) * sqrt(beamCrossSection);
        z0 = (Garfield::RndmUniform() - 0.5) * sqrt(beamCrossSection);
        break;
      }
      case BeamShape::Circle: {
        const double CrossSectionRadius = sqrt(beamCrossSection / Pi);
        x0 = Garfield::RndmUniform() * xChamberGap;
        double r = CrossSectionRadius * sqrt(Garfield::RndmUniform());
        double theta = TwoPi * Garfield::RndmUniform();
        y0 = r * cos(theta);
        z0 = r * sin(theta);
        break;
      }
    }
    electrons.push_back(P{x0, y0, z0, t, multiplicity});
    ions.push_back(P{x0, y0, z0, t, multiplicity});
  }

  // -- Simulation -- //
  // Create buffers
  const int T = omp_get_max_threads();
  std::vector<std::vector<P>> electrons_b(T);
  std::vector<std::vector<P>> ions_b(T);
  std::vector<std::vector<P>> negions_b(T);

  double t = tStartSimu;
  while (t < tEndSimu) {
    for (int k=0;k<T;++k) {
      electrons_b[k].clear();
      ions_b[k].clear();
      negions_b[k].clear();
    }
      
    // --- Drift e- ------------------------------------------------------
    const int nelectrons = electrons.size();
    #pragma omp parallel
    {
      const int tid = omp_get_thread_num();
      #pragma omp for reduction(+: nAttaElectrons)
      for (int i = 0; i < nelectrons; ++i) {
        AvalancheMicroscopic aval;
        aval.SetSensor(&sensor);
        aval.SetTimeWindow(t, t+dt);
        aval.DriftElectron(electrons[i].x, electrons[i].y, electrons[i].z,
                           electrons[i].t, 0.1, 0., 0., 0., electrons[i].w);
        for (const auto& p : aval.GetElectrons()) {
          const auto& p1 = p.path.back();
          if (p.status == -7) {
            nAttaElectrons++;
            negions_b[tid].push_back(P{p1.x, p1.y, p1.z, p1.t, p.weight});
          }
          else electrons_b[tid].push_back(P{p1.x, p1.y, p1.z, p1.t, p.weight});
        }
      }
    }

    for (auto& b : negions_b) {
      negions.insert(negions.end(), b.begin(), b.end());
    }
    for (auto& k : negions_b) k.clear();

    // --- Drift i- ------------------------------------------------------
    const int nnegions = negions.size();
    #pragma omp parallel
    {
      const int tid = omp_get_thread_num();
      #pragma omp for reduction(+: nRecombNegions)
      for (int i = 0; i < nnegions; ++i) {
        AvalancheMC driftnegion;
        driftnegion.SetSensor(&sensor);
        driftnegion.EnableRecombination(withIonRecombination, alpha);
        driftnegion.EnableDiffusion(withIonDiffusion);
        driftnegion.EnableDensityMap(withDensityMap);
        driftnegion.SetTimeSteps(.2); // [ns] default 0.02
        driftnegion.SetTimeWindow(t, t+dt);
        driftnegion.DriftNegativeIon(negions[i].x, negions[i].y, negions[i].z,
                                     negions[i].t, negions[i].w);
        for (const auto& p : driftnegion.GetNegativeIons()) {
          const auto& p1 = p.path.back();
          if (p.status == -9) nRecombNegions++; 
          else negions_b[tid].push_back(P{p1.x, p1.y, p1.z, p1.t, p.weight});
        }
      }
    }

    // --- Drift i+ ------------------------------------------------------
    const int nions = ions.size();
    #pragma omp parallel
    {
      const int tid = omp_get_thread_num();
      #pragma omp for reduction(+: nRecombIons)
      for (int i = 0; i < nions; ++i) {
        AvalancheMC driftion;
        driftion.SetSensor(&sensor);
        driftion.EnableRecombination(withIonRecombination, alpha);
        driftion.EnableDiffusion(withIonDiffusion);
        driftion.EnableDensityMap(withDensityMap);
        driftion.SetTimeSteps(.2); // [ns] default 0.02
        driftion.SetTimeWindow(t, t+dt);
        driftion.DriftIon(ions[i].x, ions[i].y, ions[i].z,
                          ions[i].t, ions[i].w);
        for (const auto& p : driftion.GetIons()) {
          const auto& p1 = p.path.back();
          if (p.status == -9) nRecombIons++;
          else ions_b[tid].push_back(P{p1.x, p1.y, p1.z, p1.t, p.weight});
        }
      }
    }
    
    // --- Update global variable and clear buffers -----------------------
    electrons.clear();
    for (auto& b : electrons_b) {
      electrons.insert(electrons.end(), b.begin(), b.end());
    }
    negions.clear();
    for (auto& b : negions_b) {
      negions.insert(negions.end(), b.begin(), b.end());
    }
    ions.clear();
    for (auto& b : ions_b) {
      ions.insert(ions.end(), b.begin(), b.end());
    }

    grid.ClearFields();
    grid.SetUniformElectricField(0., 0., 0.);
    for (auto& ion : ions) grid.AddIon(ion.x, ion.y, ion.z, ion.w);
    for (auto& nion : negions) grid.AddNegativeIon(nion.x, nion.y, nion.z, nion.w);

    t += dt;
  }

  auto end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> duration = end - start;
  auto now = std::chrono::system_clock::now();
  std::time_t currentTime = std::chrono::system_clock::to_time_t(now);
  std::tm localTime = *std::localtime(&currentTime);

  // -- Print and save results -- //
  std::ostringstream oss;

  oss << "Duration : " << duration.count() << " seconds" << std::endl;
  oss << "Number of threads used: " << T << std::endl;
  oss << "Number of pairs simulated: " << nPairs << std::endl;
  oss << "Multiplicity: " << multiplicity << std::endl;
  oss << "Ionisation density: " << densityTarget << " [cm^-3]" << std::endl;
  oss << "Number of e- attached: " << nAttaElectrons << std::endl;
  oss << "Number of i+ recombined: " << nRecombIons << std::endl;
  oss << "Number of i- recombined: " << nRecombNegions << std::endl;
  oss << "FEF: " << 1 - ((double)nAttaElectrons / nPairs) << std::endl;
  oss << "CCE: " << 1 - ((double)nRecombIons / nPairs) << std::endl;
  oss << "--------------------------------" << std::endl;
  oss << "Time : " << std::put_time(&localTime, "%H:%M:%S") << std::endl;
  oss << "Duration : " << duration.count() << " seconds" << std::endl;

  std::cout << oss.str();
  app.Run();
  return 0;
}

// OMP_NUM_THREADS=6 ./AskedDensity_mt --density 1e10