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

#include <cmath>
#include <iostream>

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

using namespace Garfield;

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

  // Gas mixture.
  MediumMagboltz gas("ar", 93., "co2", 5., "iC4H10", 2.);

  // Two regions: drift gap and amplification gap.
  ComponentAnalyticField cmpD;
  ComponentAnalyticField cmpA;
  cmpD.SetMedium(&gas);
  cmpA.SetMedium(&gas);

  // Magnetic field [Tesla].
  constexpr double bfield = 0.;
  cmpD.SetMagneticField(0., 0., bfield);
  cmpA.SetMagneticField(0., 0., bfield);

  // Amplification gap.
  constexpr double yMesh = 0.0128;   // 128 um
  constexpr double vMesh = -500.;
  cmpA.AddPlaneY(yMesh, vMesh, "mesh");
  cmpA.AddPlaneY(0., 0., "bottom");

  // Drift gap.
  constexpr double yDrift = yMesh + 0.5;  // 5 mm
  constexpr double vDrift = -1000.;
  cmpD.AddPlaneY(yDrift, vDrift, "top");
  cmpD.AddPlaneY(yMesh, vMesh, "mesh");

  // Finite anode: 1 cm x 1 cm on the plane y = 0.
  constexpr double xAnodeMin = -1;  // cm
  constexpr double xAnodeMax =  1;  // cm
  constexpr double zAnodeMin = -1;  // cm
  constexpr double zAnodeMax =  1;  // cm

  cmpA.AddPixelOnPlaneY(0., xAnodeMin, xAnodeMax,
                        zAnodeMin, zAnodeMax, "anode");

  // Sensor.
  Sensor sensor;
  sensor.AddComponent(&cmpD);
  sensor.AddComponent(&cmpA);
  sensor.SetArea(xAnodeMin, 0., zAnodeMin,
                 xAnodeMax, yDrift, zAnodeMax);

  // Request signal calculation on the anode.
  sensor.AddElectrode(&cmpA, "anode");

  // Time window [ns].
  const double tMin = 0.;
  const double tMax = 100.;
  const double tStep = 0.05;
  const int nTimeBins = int((tMax - tMin) / tStep);
  sensor.SetTimeWindow(tMin, tStep, nTimeBins);

  // Microscopic avalanche.
  AvalancheMicroscopic aval(&sensor);

  // Heed track.
  TrackHeed track(&sensor);
  track.SetParticle("muon");
  constexpr double momentum = 50.e9;  // [eV / c]
  track.SetMomentum(momentum);
  track.EnableMagneticField();

  // Plotting.
  ViewDrift driftView;
  track.EnablePlotting(&driftView);
  aval.EnablePlotting(&driftView);
  aval.EnableExcitationMarkers(false);
  aval.EnableIonisationMarkers(false);

  // Incoming particle start point.
  const double xt = 0.0;      // cm
  const double yt = yDrift;   // cm
  const double zt = 0.0;      // cm

  // Counters.
  int nPrimary = 0;
  int nReachedMesh = 0;
  long long nCollected = 0;

  // Elementary charge in fC.
  constexpr double eCharge = 1.602176634e-4;

  // Simulate the track.
  track.NewTrack(xt, yt, zt, 0., 0., -1., 0.);

  // Loop over Heed clusters.
  for (const auto& cluster : track.GetClusters()) {
    for (const auto& electron : cluster.electrons) {
      ++nPrimary;

      // Step 1: drift in the drift gap up to the mesh.
      aval.AvalancheElectron(electron.x, electron.y, electron.z, electron.t,
                             3, 0., 0., 0.);

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

      // Endpoint of the initial electron in the drift region.
      const auto& p1 = driftElectrons.front().path.back();

      // Only continue if the electron reached the mesh plane.
      if (std::fabs(p1.y - yMesh) > 1.e-6) continue;
      ++nReachedMesh;

      // Step 2: move electron just below the mesh into amplification gap.
      aval.AvalancheElectron(p1.x, yMesh - 1.e-6, p1.z, p1.t,
                             p1.energy, 0., -1., 0.);

      // Count how many avalanche electrons actually end on the anode plane.
      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) < 1.e-6 &&
            pend.x >= xAnodeMin && pend.x <= xAnodeMax &&
            pend.z >= zAnodeMin && pend.z <= zAnodeMax) {
          ++nCollected;
        }
      }
    }
  }

  // Integrated signal on the anode.
  sensor.IntegrateSignals();
  const double q = sensor.GetSignal("anode", nTimeBins - 1);
  const double nEq = std::fabs(q) / eCharge;

  std::cout << "---------------------------\n";
  std::cout << "Primary electrons from Heed:     " << nPrimary << "\n";
  std::cout << "Primary electrons reaching mesh: " << nReachedMesh << "\n";
  std::cout << "Integrated anode signal Q [fC]:  " << q << "\n";
  std::cout << "Equivalent electrons |Q|/e:      " << nEq << "\n";
  std::cout << "Collected electrons by endpoint: " << nCollected << "\n";
  std::cout << "---------------------------\n";

  // Plotting.
  gStyle->SetPadRightMargin(0.15);
  TCanvas* c1 = new TCanvas("c1", "", 1400, 600);
  c1->Divide(2, 1);

  driftView.SetArea(xAnodeMin, 0.0, xAnodeMax, yDrift);
  driftView.SetClusterMarkerSize(0.1);
  driftView.SetCollisionMarkerSize(0.5);
  driftView.SetCanvas((TPad*)c1->cd(1));
  driftView.Plot(true);

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

  c1->SaveAs("plot.pdf");
  app.Run();
  return 0;
}