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

#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/ViewSignal.hh"
#include "Garfield/FundamentalConstants.hh"

using namespace Garfield;
using namespace std;

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");
  cmpA.AddReadout("pad");
  // 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");
  cmpD.AddReadout("drift");
  // 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(&cmpA, "pad"); 
  sensor.AddElectrode(&cmpD, "mesh2");
  sensor.AddElectrode(&cmpD, "drift");

  // Set the time window [ns] for the signal calculation.
  const double tMin = 0.; 
  // const double tMax = 220.; 
  const double tMax = 480.; 
  const double tStep = 0.05; 
  // const double tMax = 500000.; 
  // const double tStep = 5.; 
  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);
  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();
  drift.UseWeightingPotential(true);

  //track.EnableTransverseStraggling(false); //false=polla swmatidia, true=1
  ViewSignal signalView;
  constexpr bool plotSignal = true;
  if (plotSignal) {
    signalView.SetSensor(&sensor);
  } 

  TCanvas* ct = nullptr;
  ViewDrift driftView;
  constexpr bool plotDrift = false;
  if (plotDrift) {
    ct = new TCanvas("ct", "", 600, 600);
    drift.EnablePlotting(&driftView);
    track.EnablePlotting(&driftView);
  }

  const unsigned int nTracks = 1;
  for (unsigned int j = 0; j < nTracks; ++j) {
    sensor.ClearSignal();
    track.NewTrack(0., yDrift - 1.e-10, 0., 0., 0., -0.6, 0. );
    int nsum = 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)) {
      nsum += nc;
      drift.SetIonSignalScalingFactor(nc);
      drift.DriftIon(xc, yc, zc, tc);
      drift.SetElectronSignalScalingFactor(nc);
      drift.DriftElectron(xc, yc, zc, tc);
      double xe = 0., ye = 0., ze = 0., te = 0.;
      int status = 0;
      drift.GetEndPoint(xe, ye, ze, te, status);
      ye = yMesh - 1.e-10;
      drift.DriftElectron(xe, ye, ze, te);
    }
    sensor.IntegrateSignals();
    if (plotSignal) signalView.PlotSignal("mesh2");
  
    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); 
    const double qpad = sensor.GetSignal("pad", nTimeBins - 1);
    const double qdrift = sensor.GetSignal("drift", nTimeBins - 1);
 
    std::cout << "  mesh2 tot=  " << qmesh2 << "  mesh2 e=  " << qemesh2 << "  mesh2 ion=  " << ionmesh2 << "  mesh1 tot " << qmesh1 << "  mesh1 e " << qemesh1 << "   mesh1 ion  " << ionmesh1 << "\n";

    std::cout << "Deposited charge: " << nsum * ElementaryCharge << "\n";
  }

  if (plotDrift) {
    driftView.SetArea(-6, 0., 6, yDrift);
    driftView.Plot2d(true);
  }

  app.Run(true);
}