#include <iostream>
#include <fstream>
#include <omp.h>

#include <TCanvas.h>
#include <TROOT.h>
#include <TApplication.h>

#include "Garfield/AvalancheMC.hh"
#include "Garfield/DriftLineRKF.hh"
#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Plotting.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/TrackSrim.hh"
#include "Garfield/ViewSignal.hh"

using namespace Garfield;

int main(int argc, char * argv[]) {
  TApplication app("app", &argc, argv);

  MediumMagboltz gas;
  if (!gas.LoadGasFile("Ar_80_CO2_20.gas") || !gas.LoadIonMobility("IonMobility_Ar+_Ar.txt"))
    return EXIT_FAILURE;

  ComponentAnalyticField cmp;
  cmp.SetMedium(&gas);
  const double rWire = 15.e-4;
  const double rTube = 3.5;
  const double vWire = 1100.;
  const double vTube = 0.;
  // Add the wire in the centre.
  cmp.AddWire(0., 0., 2. * rWire, vWire, "s");
  // Add the tube.
  cmp.AddTube(rTube, vTube, 0);

  Sensor sensor(&cmp);
  sensor.AddElectrode(&cmp, "s");
  // Set the signal time window.
  const double tstep = 0.2;
  const double tmin = 4000.;
  const double tmax = 10000.;
  const unsigned int nbins = (tmax - tmin) / tstep;
  sensor.SetTimeWindow(tmin, tstep, nbins);

  TrackSrim tr(&sensor);
  const std::string file = "Li_in_Ar_CO2_gas.txt";
  if (!tr.ReadFile(file)) {
    std::cerr << "Reading SRIM file failed.\n";
  }
  tr.SetKineticEnergy(0.84e6);
  // tr.SetTargetClusterSize(10);
  tr.EnableTransverseStraggling(false);
  tr.EnableLongitudinalStraggling(false);
  tr.Print();

  const double rTrack = 0.3;
  const double projectiley = -sqrt(rTube * rTube - rTrack * rTrack);
  if (!tr.NewTrack(rTrack, projectiley, 0., 0., 0., 1., 0.)) {
    std::cerr << "Generating clusters failed.\n";
    return 0;
  }
  unsigned int nn = 0;
  unsigned int nclusters = tr.GetClusters().size();
  std::cout << nclusters << " clusters on the track.\n";
  for (const auto& cluster : tr.GetClusters()) {
    std::cout << "Cluster with " << cluster.n << " electrons" << std::endl;
    #pragma omp parallel for schedule(dynamic)
    for (int k = 0; k < cluster.n; ++k) {
      
      AvalancheMC driftelectron(&sensor);
      AvalancheMC driftion(&sensor);
      driftelectron.EnableRKFSteps(true);
      driftelectron.SetDistanceSteps(0.01);
      driftion.EnableRKFSteps(true);
      driftion.SetDistanceSteps(0.01);
      driftelectron.AvalancheElectron(cluster.x, cluster.y, cluster.z, cluster.t);
      // Drift the ions created in the avalanche.
      for (const auto& electron : driftelectron.GetElectrons()) {
        const auto& p0 = electron.path[0];
        driftion.DriftIon(p0.x, p0.y, p0.z, p0.t);       
      }
      /*
      DriftLineRKF driftRKF(&sensor);
      driftRKF.EnableSignalCalculation(true);
      driftRKF.EnableIonTail(true);
      driftRKF.SetMaximumStepSize(1E-2);
      driftRKF.SetIntegrationAccuracy(1E-5);
      driftRKF.DriftElectron(cluster.x, cluster.y, cluster.z, cluster.t);
      */
    }
    std::cout << ++nn << "/" << nclusters << " clusters" << std::endl;
  }

  TCanvas* cS = new TCanvas("cS", "", 600, 600);
  ViewSignal signalView(&sensor);
  signalView.SetCanvas(cS);
  signalView.PlotSignal("s", "tei");
  std::cout << "Calculation is finished.\n";
  app.Run();
}