#include <iostream>
#include <fstream>
#include <cstdlib>

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

#include "Garfield/ViewCell.hh"
#include "Garfield/ViewDrift.hh"
#include "Garfield/ViewSignal.hh"
#include "Garfield/ViewField.hh"

#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/DriftLineRKF.hh"
#include "Garfield/AvalancheMicroscopic.hh"
#include "Garfield/AvalancheMC.hh"

using namespace Garfield;

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

  TApplication app("app", &argc, argv);
  MediumMagboltz gas;
  // Read the electron transport coefficients from a .gas file.
  gas.LoadGasFile("ar_c4h10.gas");
  auto installdir = std::getenv("GARFIELD_INSTALL");
  if (!installdir) {
    std::cerr << "GARFIELD_INSTALL variable not set.\n";
    return 1;
  }
  const std::string path = installdir;
  gas.LoadIonMobility(path + "/share/Garfield/Data/IonMobility_Ar+_Ar.txt");

  // Make a component with analytic electric field.
  ComponentAnalyticField cmp;
  cmp.SetMedium(&gas);   
  // Distance between grid and anode wires [cm].
  constexpr double gap = 0.4;
  // Grid wire spacing (0.5 mm).
  constexpr double periodg = 0.05;
  // Grid wire diameter (50 um).
  constexpr double dg = 0.0050;
   // Anode diameter (20 um).
  constexpr double da = 0.0020;

  // Grid wire potential (ground).
  constexpr double vg = 0.;         
  // Anode wire potential.
  constexpr double va = 1500.;

  // Periodicity (anode wire spacing) [cm].
  constexpr double perioda = 0.2;    
  cmp.SetPeriodicityX(perioda); 

  // Add the grid wires.
  constexpr double yg = 1 * gap;
  constexpr double xg0 = 0 * periodg;
  constexpr double xg1 = 1 * periodg;
  constexpr double xg2 = 2 * periodg;
  constexpr double xg3 = 3 * periodg;
  cmp.AddWire(xg0, yg, dg, vg, "g");
  cmp.AddWire(xg1, yg, dg, vg, "g");
  cmp.AddWire(xg2, yg, dg, vg, "g");
  cmp.AddWire(xg3, yg, dg, vg, "g");

  // Add the anode wires. 
  constexpr double xa = 0.;
  constexpr double ya = 0.;
  cmp.AddWire(xa, ya, da, va, "a");

 // Add the cathode and pad planes.
  constexpr double ycathode = 3.4;
  constexpr double vcathode = -390.;
  constexpr double ypad = -2.6;
  constexpr double vpad = 0; // pad is grounded
  cmp.AddPlaneY(ycathode, vcathode, "cathode");
  cmp.AddPlaneY(ypad, vpad, "pad");

  // Request calculation of the weighting field. 
  cmp.AddReadout("a");      

  const double xmin = -10 * periodg; 
  const double xmax =  10 * periodg;

  // Make a sensor.
  Sensor sensor;
  sensor.AddComponent(&cmp);
  sensor.SetArea(xmin, ypad, -1., xmax, ycathode, 1.);
  sensor.AddElectrode(&cmp, "a");
   
  /*
  // Draw equipotential lines.
  ViewField fieldView;
  fieldView.SetComponent(&cmp);
  fieldView.SetArea(xmin, -0.2, xmax, 0.2);
  fieldView.SetVoltageRange(-600., 2000.);
  fieldView.SetNumberOfContours(100);
  fieldView.PlotContour();
  */

  // Set the time window and binning for the signal calculation.
  const double tstep = 0.5;
  const double tmin = 0. * tstep;
  const unsigned int nbins = 12000;
  sensor.SetTimeWindow(0, tstep, nbins);

  DriftLineRKF drift;
  drift.SetSensor(&sensor);
  drift.SetMaximumStepSize(5.e-4);
  drift.EnableIonTail();
  drift.EnableSignalCalculation();

  AvalancheMicroscopic aval;
  aval.SetSensor(&sensor);
  // aval.EnableAvalancheSizeLimit(3000.);   
  aval.EnableSignalCalculation();

  AvalancheMC mc;
  mc.SetSensor(&sensor);
  mc.SetDistanceSteps(2.e-4);
  mc.EnableSignalCalculation();

  TCanvas* cD = new TCanvas("cD", "", 600, 600);
  ViewCell cellView;
  ViewDrift driftView;
  constexpr bool plotDrift = true;
  if (plotDrift) {
    cellView.SetCanvas(cD);
    cellView.SetComponent(&cmp);
    driftView.SetArea(xmin, ypad, xmax, ycathode);
    driftView.SetCanvas(cD);         
    drift.EnablePlotting(&driftView);
    aval.EnablePlotting(&driftView);
    mc.EnablePlotting(&driftView);
  }

  constexpr unsigned int nEvents = 1;
  for (unsigned int i = 0; i < nEvents; ++i) { 
    std::cout << i << "/" << nEvents << "\n";
    sensor.ClearSignal();
    const double x0 = 0.025;
    const double y0 = 2;
    const double z0 = 0;
    const double t0 = 0;
    const double e0 = 0.1;

    aval.AvalancheElectron(x0, y0, z0, t0, e0, 0., 0., 0.);
    // mc.AvalancheElectron(x0, y0, z0, t0);
    double xe1 = 0., ye1 = 0., ze1 = 0., te1 = 0., e1 = 0.;
    int status = 0;
    double xe2 = 0., ye2 = 0., ze2 = 0., te2 = 0., e2 = 0.;          
    const unsigned int np = aval.GetNumberOfElectronEndpoints();
    std::cout << np << " electron trajectories.\n";
    for (unsigned int j = 0; j < np; ++j) {
      aval.GetElectronEndpoint(j, xe1, ye1, ze1, te1, e1, 
                                  xe2, ye2, ze2, te2, e2, status);
      mc.DriftIon(xe1, ye1, ze1, te1);
    }
  }

  if (plotDrift) {
    cD->Clear();
    cellView.SetArea(xmin, ypad, xmax, ycathode);
    driftView.SetArea(xmin, ypad, xmax, ycathode);
    constexpr bool twod = true;
    constexpr bool drawaxis = false;     
    driftView.Plot(twod, drawaxis);
    cellView.Plot2d();           
  }

  std::cout << "Finished.\n";
  app.Run(true);
}