#include <cmath>
#include <iostream>

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

#include "Garfield/AvalancheMicroscopic.hh"
#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Random.hh"
#include "Garfield/RandomEngineRoot.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/ViewDrift.hh"
#include "Garfield/ViewField.hh"

using namespace Garfield;

int main(int argc, char* argv[]) {
  TApplication app("app", &argc, argv);
  std::ofstream fout("anode_hits.txt");
  fout << "# event x[cm] y[cm] z[cm]\n";

  // Random seed for reproducibility.
  RandomEngineRoot randomEngine(123456);
  Random::SetEngine(randomEngine);

  // --------------------------------------------------
  // Gas
  // --------------------------------------------------
  MediumMagboltz gas("ar", 93., "co2", 5., "ic4h10", 2.);
  gas.SetTemperature(293.15);
  gas.SetPressure(760.);
gas.LoadGasFile("ar_93_co2_5_ic4h10_2.gas");   
// Penning transfer probability.
constexpr double rPenning = 0.42;
// Mean distance from the point of excitation.
constexpr double lambdaPenning = 0.;
gas.EnablePenningTransfer(rPenning, lambdaPenning, "ar");
 // --------------------------------------------------
  // Geometry: drift region + amplification region
  // --------------------------------------------------
  ComponentAnalyticField cmpD;
  ComponentAnalyticField cmpA;
  cmpD.SetMedium(&gas);
  cmpA.SetMedium(&gas);

  constexpr double bfield = 0.;
  cmpD.SetMagneticField(0., 0., bfield);
  cmpA.SetMagneticField(0., 0., bfield);

  // Amplification gap [cm]
  constexpr double yMesh = 0.0128;
  constexpr double vMesh = 500.;
  cmpA.AddPlaneY(yMesh, 0, "mesh");
  cmpA.AddPlaneY(0., 500., "bottom");

  // Drift gap [cm]
  constexpr double yDrift = yMesh + 0.7;
  constexpr double vDrift = -300.;
  cmpD.AddPlaneY(yDrift, vDrift, "top");
  cmpD.AddPlaneY(yMesh, 0, "mesh");

  // One strip on the anode plane.
  constexpr double pitch = 3.;
  cmpA.AddStripOnPlaneY('z', 0., -pitch, pitch, "strip1");

  // --------------------------------------------------
  // Sensor
  // --------------------------------------------------
  Sensor sensor;
  sensor.AddComponent(&cmpD);
  sensor.AddComponent(&cmpA);
  sensor.SetArea(-pitch, 0., -3., pitch, yDrift, 3.);
  sensor.AddElectrode(&cmpA, "strip1");

  // --------------------------------------------------
  // Avalanche object for statistics
  // --------------------------------------------------
  AvalancheMicroscopic aval(&sensor);
  aval.EnableExcitationMarkers(false);
  aval.EnableIonisationMarkers(false);

  // --------------------------------------------------
  // Parameters of the primary electrons
  // --------------------------------------------------
  constexpr int nPrimaryElectrons = 100;

  // Start slightly above the mesh, in the drift region.
  constexpr double dyStart = 1.e-3;   // [cm] = 10 micron above mesh
  const double yStart = yMesh + dyStart;

  // Fixed starting position.
  const double xStart = 0.0;    // [cm]
  const double zStart = 0.0;    // [cm]
  const double tStart = 0.0;    // [ns]
  const double eStart = 0.1;    // [eV], small initial energy

  // Counters
  int nReachedMesh = 0;
  int nTransferred = 0;
  long long totalAvalancheElectrons = 0;
  long long totalAvalancheIons = 0;
  long long totalElectronsOnAnode = 0;

  const double tolYAnode = 1.e-6;

  // --------------------------------------------------
  // Simulate primary electrons
  // --------------------------------------------------
  for (int i = 1; i <= nPrimaryElectrons; ++i) {
    if (i % 1 == 0) {
      std::cout << i << " / " << nPrimaryElectrons << std::endl;
    }

    // Stage 1:
    // drift the electron from just above the mesh down to the mesh.
    aval.AvalancheElectron(xStart, yStart, zStart, tStart,
                           eStart, 0., -1., 0.);

    const auto& electrons1 = aval.GetElectrons();
    if (electrons1.empty()) continue;
    if (electrons1.front().path.empty()) continue;

    // Endpoint of the initial electron track.
    const auto& p1 = electrons1.front().path.back();

    // We only keep electrons that really reached the mesh plane.
    if (std::fabs(p1.y - yMesh) > 1.e-6) continue;
    ++nReachedMesh;

    // Stage 2:
    // move the electron slightly below the mesh and simulate avalanche.
    aval.AvalancheElectron(p1.x, yMesh - 1.e-6, p1.z, p1.t,
                           p1.energy, 0., -1., 0.);
    ++nTransferred;

    int ne = 0, ni = 0;
    aval.GetAvalancheSize(ne, ni);
    totalAvalancheElectrons += ne;
    totalAvalancheIons += ni;

    // Count electrons that end on the anode plane.
    int nOnAnodeThisEvent = 0;
    const auto& avalElectrons = aval.GetElectrons();

    for (const auto& el : avalElectrons) {
      if (el.path.empty()) continue;

      const auto& pend = el.path.back();

      if (std::fabs(pend.y - 0.0) <= tolYAnode &&
          pend.z >= -pitch && pend.z <= pitch) {
      
        ++nOnAnodeThisEvent;
      
        fout << i << " "
             << pend.x << " "
             << pend.y << " "
             << pend.z << "\n";
      
      }
    }

    totalElectronsOnAnode += nOnAnodeThisEvent;
  }

  // --------------------------------------------------
  // Print summary
  // --------------------------------------------------
  std::cout << "\n========== Summary ==========\n";
  std::cout << "Primary electrons started: " << nPrimaryElectrons << "\n";
  std::cout << "Reached mesh plane:        " << nReachedMesh << "\n";
  std::cout << "Transferred below mesh:    " << nTransferred << "\n";

  if (nTransferred > 0) {
    std::cout << "Mean avalanche electrons:  "
              << double(totalAvalancheElectrons) / nTransferred << "\n";
    std::cout << "Mean avalanche ions:       "
              << double(totalAvalancheIons) / nTransferred << "\n";
    std::cout << "Mean electrons on anode:   "
              << double(totalElectronsOnAnode) / nTransferred << "\n";
  }
  std::cout << "=============================\n";

  // --------------------------------------------------
  // Separate object for plotting only the amplification gap
  // --------------------------------------------------
  ViewDrift driftView;
  driftView.SetArea(-0.01, 0.0, 0.01, yMesh);
  driftView.SetClusterMarkerSize(0.1);
  driftView.SetCollisionMarkerSize(0.5);

  AvalancheMicroscopic avalPlot(&sensor);
  avalPlot.EnablePlotting(&driftView);
  avalPlot.EnableExcitationMarkers(false);
  avalPlot.EnableIonisationMarkers(false);

  constexpr int nPlotElectrons = 1;
  for (int i = 1; i <= nPlotElectrons; ++i) {
    // Start directly in the amplification gap for plotting only.
    avalPlot.AvalancheElectron(xStart, yMesh - 1.e-6, zStart, tStart,
                               eStart, 0., -1., 0.);
  }

  // --------------------------------------------------
  // Plot drift lines + potential (only amplification gap)
  // --------------------------------------------------
  gStyle->SetPadRightMargin(0.15);
  TCanvas* c1 = new TCanvas("c1", "", 1400, 600);
  c1->Divide(2, 1);

  driftView.SetCanvas((TPad*)c1->cd(1));
  driftView.Plot(true);

  ViewField fieldView(&sensor);
  fieldView.SetCanvas((TPad*)c1->cd(2));
  fieldView.SetArea(-0.2, 0.0, 0.2, yMesh);
  fieldView.Plot("v", "cont1z");

  c1->SaveAs("plot.pdf");

  app.Run();
  return 0;
}