#include <iostream>
#include <fstream>
#include <cstdlib>
#include <cmath>
#include <random>

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

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

using namespace Garfield;

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

  TApplication app("app", &argc, argv);

  // Set the gas mixture.
  MediumMagboltz gas;
  // gas.LoadGasFile("donegas.gas");
  gas.LoadGasFile("ar_93_co2_7.gas");
  const std::string path = std::getenv("GARFIELD_HOME");
  gas.LoadIonMobility(path + "/Data/IonMobility_Ar+_Ar.txt");

  // Build the Micromegas geometry.
  // We use separate components for drift and amplification regions. 

  // 50 um amplification gap. 
  ComponentAnalyticField cmpA; 
  cmpA.SetMedium(&gas);
  constexpr double yMesh = 0.005;
  constexpr double vMesh = -300.;
  cmpA.AddPlaneY(yMesh, vMesh, "mesh1");
  cmpA.AddPlaneY(0., 0., "pad");
  cmpA.AddReadout("mesh1");

  // 6 mm drift gap.
  ComponentAnalyticField cmpD;
  cmpD.SetMedium(&gas);
  constexpr double yDrift = yMesh + 0.6; 
  constexpr double vDrift = -800.; 
  cmpD.AddPlaneY(yDrift, vDrift, "drift");
  cmpD.AddPlaneY(yMesh, vMesh, "mesh2");
  cmpD.AddReadout("mesh2");

  // Assemble a sensor.
  Sensor sensor;
  sensor.AddComponent(&cmpA); 
  sensor.AddComponent(&cmpD); 
  sensor.SetArea(-4.75, 0., -4.75, 4.75, yDrift, 4.75);
  sensor.AddElectrode(&cmpA, "mesh1"); 
  sensor.AddElectrode(&cmpD, "mesh2");
  sensor.ClearSignal();

  // Set the time window [ns] for the signal calculation.
  const double tMin = 0.; 
  const double tMax = 220.; 
  const double tStep = 0.05; 
  const int nTimeBins = int((tMax - tMin) / tStep); 
  sensor.SetTimeWindow(0., tStep, nTimeBins);

  TrackSrim track;
  track.SetSensor(&sensor);
  // Read SRIM output from file.
  const std::string file = "alpha_in_93ar-7co2_0.01-5mev.txt";
  if (!track.ReadFile(file)) {
    std::cerr << "Reading SRIM file failed.\n";
    return 0;
  }
  // Set the initial kinetic energy of the particle (in eV).
  track.SetKineticEnergy(5e6);
  // Set the W value and Fano factor of the gas.
  track.SetWorkFunction(26.0);
  track.SetFanoFactor(0.3);
  // Set A and Z of the gas (not sure what's the correct mixing law).
  const double za = 0.93 * (18. / 40.) + 0.07 * (22. / 44.);
  track.SetAtomicMassNumbers(22. / za, 22);
  track.SetTargetClusterSize(500); 

  /*
  // RKF integration.
  DriftLineRKF drift;
  drift.SetSensor(&sensor);
  drift.SetMaximumStepSize(5.e-4);
  //drift.SetGainFluctuationsPolya(0., 20000.);
  drift.EnableIonTail();
  drift.EnableSignalCalculation();
  */

  AvalancheMicroscopic aval;
  aval.SetSensor(&sensor);
  aval.EnableSignalCalculation();

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

  ViewDrift driftView;
  constexpr bool plotDrift = false;
  if (plotDrift) {
    //drift.EnablePlotting(&driftView);
    aval.EnablePlotting(&driftView);
    mc.EnablePlotting(&driftView);
    track.EnablePlotting(&driftView);
  }

  std::ofstream outfile1;
  outfile1.open("0_kathetoaval.txt");

  const unsigned int nTracks = 10;
  for (unsigned int j = 0; j < nTracks; ++j) {
    std::cout << "Track " << j << "...\n";
    sensor.ClearSignal();
    track.NewTrack(0, yDrift - 1.e-10, 0, 0, 0, -1., 0);
    double xc = 0., yc = 0., zc = 0., tc = 0., ec = 0., extra = 0.;
    int nc = 0;
    while (track.GetCluster(xc, yc, zc, tc, nc, ec, extra)) {   
      for (int k = 0; k < nc; ++k) {
        mc.DriftIon(xc, yc, zc, tc);
        double xe = 0., ye = 0., ze = 0., te = 0., ee = 0.1;
        if (yc > yMesh) {
          // Cluster in the drift gap.
          aval.DriftElectron(xc, yc, zc, tc, 0.1, 0., 0., 0.);
          // Get the electron's end point.
          double x1 = 0., y1 = 0., z1 = 0., t1 = 0., e1 = 0.;
          int status = 0;
          aval.GetElectronEndpoint(0, x1, y1, z1, t1, e1, 
                                      xe, ye, ze, te, ee, status);
          if (status != -11 || fabs(ye - yMesh) > 1.e-6) {
            // Skip electrons that didn't end up at the mesh plane.
            continue;
          }
          // Move the electron from the drift to the avalanche gap.
          ye = yMesh - 1.e-6;
        } else {
          // Cluster in the avalanche gap.
          xe = xc;
          ye = yc;
          ze = zc;
          te = tc;
        }
        // Simulate the electron avalanche in the multiplication gap.
        aval.AvalancheElectron(xe, ye, ze, te, ee, 0., 0., 0.);
        // Loop over the electrons in the avalanche.
        const unsigned int np = aval.GetNumberOfElectronEndpoints();
        for (unsigned int i = 0; i < np; ++i) {
          double x1 = 0., y1 = 0., z1 = 0., t1 = 0., e1 = 0.;
          double x2 = 0., y2 = 0., z2 = 0., t2 = 0., e2 = 0.;
          int status = 0;
          aval.GetElectronEndpoint(i, x1, y1, z1, t1, e1,
                                      x2, y2, z2, t2, e2, status);
          // Skip the initial electron.
          if (fabs(y1 - yMesh) < 1.1e-6) continue;
          // Simulate the ion drift.
          mc.DriftIon(x1, y1, z1, t1);
        }
      }
    }
    sensor.IntegrateSignals();
    const double qmesh1 = sensor.GetSignal("mesh1", nTimeBins - 1);  
    const double qmesh2 = sensor.GetSignal("mesh2", nTimeBins - 1); 
    const double qemesh1 = sensor.GetElectronSignal("mesh1", nTimeBins - 1);
    const double qemesh2 = sensor.GetElectronSignal("mesh2", nTimeBins - 1); 
    const double ionmesh1 = sensor.GetIonSignal("mesh1", nTimeBins - 1); 
    const double ionmesh2 = sensor.GetIonSignal("mesh2", nTimeBins - 1); 
    outfile1 << qmesh1 << " " << qmesh2 << " " << qemesh1 << " " <<  qemesh2 << "\n";
    std::cout << "  mesh2 tot=  " << qmesh2 << "  mesh2 e=  " << qemesh2 << "  mesh2 ion=  " << ionmesh2 << "  mesh1 tot " << qmesh1 << "  mesh1 e " << qemesh1 << "   mesh1 ion  " << ionmesh1 << "\n";
  }
  outfile1.close();
  if (plotDrift) {
    driftView.SetArea(-6, 0., 6, yDrift);
    driftView.Plot2d(true);
  }

  app.Run(true);
}