#include <cmath>
#include <cmath>
#include <fstream>
#include <iostream>
#include <string>

#include "Garfield/AvalancheMC.hh"
#include "Garfield/ComponentNeBem3d.hh"
#include "Garfield/GeometrySimple.hh"
#include "Garfield/MediumConductor.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Random.hh"
#include "Garfield/RandomEngineRoot.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/SolidBox.hh"

using namespace Garfield;

int main() {
  // --------------------------------------------------
  // 1. Random engine
  // --------------------------------------------------
  RandomEngineRoot randomEngine(123456);
  Random::SetEngine(randomEngine);

  // --------------------------------------------------
  // 2. Gas + transport parameter table
  // --------------------------------------------------
  MediumMagboltz gas("ar", 93., "co2", 5., "iC4H10", 2.);
  gas.SetTemperature(293.15);
  gas.SetPressure(760.);

  const std::string gasFile = "ar_93_co2_5_ic4h10_2.gas";

  {
    std::ifstream fin(gasFile);
    if (fin.good()) {
      fin.close();
      gas.LoadGasFile(gasFile);
      std::cout << "Loaded gas table from " << gasFile << "\n";
    } else {
      std::cout << "Gas file not found, running Magboltz...\n";
      gas.SetFieldGrid(100., 1.e5, 20, true);   // 100 V/cm ... 100 kV/cm
      gas.GenerateGasTable(10, true);           // nColl = 10
      gas.WriteGasFile(gasFile);
      std::cout << "Gas table saved to " << gasFile << "\n";
    }
  }

  // --------------------------------------------------
  // 3. Conductor
  // --------------------------------------------------
  MediumConductor metal;

  // --------------------------------------------------
  // 4. Geometry
  // --------------------------------------------------
  GeometrySimple geo;
  geo.SetMedium(&gas);

  const double xCathPlane  =  0.0;
  const double xMeshPlane  = -0.5;
  const double xAnodePlane = -0.5128;

  const double yMin = 0.0;
  const double yMax = 10.0;
  const double zMin = -10.0;
  const double zMax = 0.0;

  const double hole     = 0.007;   // 70 um = 0.007 cm
  const double halfHole = hole / 2.0;

  const double yHoleCenter = 5.0;
  const double zHoleCenter = -5.0;

  const double yHoleMin = yHoleCenter - halfHole;
  const double yHoleMax = yHoleCenter + halfHole;
  const double zHoleMin = zHoleCenter - halfHole;
  const double zHoleMax = zHoleCenter + halfHole;

  const double tCath  = 0.002;
  const double tMesh  = 0.0005;
  const double tAnode = 0.002;

  const double vCath  = -300.0;
  const double vMesh  =    0.0;
  const double vAnode = +500.0;

  // Cathode
  SolidBox cathode(
      xCathPlane + tCath / 2.0, 5.0, -5.0,
      tCath / 2.0, 5.0, 5.0
  );
  cathode.SetBoundaryPotential(vCath);
  cathode.SetLabel("cathode");
  geo.AddSolid(&cathode, &metal);

  // Anode
  SolidBox anode(
      xAnodePlane - tAnode / 2.0, 5.0, -5.0,
      tAnode / 2.0, 5.0, 5.0
  );
  anode.SetBoundaryPotential(vAnode);
  anode.SetLabel("anode");
  geo.AddSolid(&anode, &metal);

  // Mesh: 4 rectangles around the hole
  SolidBox meshLeft(
      xMeshPlane,
      0.5 * (yMin + yHoleMin),
      0.5 * (zMin + zMax),
      tMesh / 2.0,
      0.5 * (yHoleMin - yMin),
      0.5 * (zMax - zMin)
  );
  meshLeft.SetBoundaryPotential(vMesh);
  meshLeft.SetLabel("mesh");
  geo.AddSolid(&meshLeft, &metal);

  SolidBox meshRight(
      xMeshPlane,
      0.5 * (yHoleMax + yMax),
      0.5 * (zMin + zMax),
      tMesh / 2.0,
      0.5 * (yMax - yHoleMax),
      0.5 * (zMax - zMin)
  );
  meshRight.SetBoundaryPotential(vMesh);
  meshRight.SetLabel("mesh");
  geo.AddSolid(&meshRight, &metal);

  SolidBox meshBottom(
      xMeshPlane,
      yHoleCenter,
      0.5 * (zMin + zHoleMin),
      tMesh / 2.0,
      halfHole,
      0.5 * (zHoleMin - zMin)
  );
  meshBottom.SetBoundaryPotential(vMesh);
  meshBottom.SetLabel("mesh");
  geo.AddSolid(&meshBottom, &metal);

  SolidBox meshTop(
      xMeshPlane,
      yHoleCenter,
      0.5 * (zHoleMax + zMax),
      tMesh / 2.0,
      halfHole,
      0.5 * (zMax - zHoleMax)
  );
  meshTop.SetBoundaryPotential(vMesh);
  meshTop.SetLabel("mesh");
  geo.AddSolid(&meshTop, &metal);

  // --------------------------------------------------
  // 5. neBEM field
  // --------------------------------------------------
  ComponentNeBem3d cmp;
  cmp.SetGeometry(&geo);
  cmp.SetTargetElementSize(5.e-4);
  cmp.SetMinMaxNumberOfElements(3, 7);
  cmp.UseLUInversion();
  cmp.Initialise();

  // --------------------------------------------------
  // 6. Sensor
  // --------------------------------------------------
  Sensor sensor;
  sensor.AddComponent(&cmp);
  sensor.SetArea(xAnodePlane - 0.005, yMin, zMin,
                 xCathPlane  + 0.005, yMax, zMax);

  // --------------------------------------------------
  // 7. AvalancheMC
  // --------------------------------------------------
  AvalancheMC aval(&sensor);
  aval.SetDistanceSteps(1.e-4);   // 1 um step
  aval.EnableRKFSteps(true);

  // --------------------------------------------------
  // 8. Helpers
  // --------------------------------------------------
  const int    nPrim = 10;
  const double eps   = 1.e-6;

  auto insideHoleYZ = [&](const double y, const double z) {
    return (y >= yHoleMin && y <= yHoleMax &&
            z >= zHoleMin && z <= zHoleMax);
  };

  auto nearMeshX = [&](const double x) {
    return std::abs(x - xMeshPlane) <= (0.5 * tMesh + 5.e-4);
  };

  // --------------------------------------------------
  // 9. Counters
  // --------------------------------------------------
  long long sumNeTot            = 0;
  long long sumCollectedOnAnode = 0;
  long long sumLostOnMesh       = 0;
  long long sumLostOnCathode    = 0;
  long long sumOtherEndpoints   = 0;
  int       nPrimaryCrossedMesh = 0;

  std::ofstream out("electron_endpoints.txt");
  out << "# event  electron  x1 y1 z1  x2 y2 z2  status\n";

  // --------------------------------------------------
  // 10. Loop over primary electrons
  // --------------------------------------------------
  for (int i = 0; i < nPrim; ++i) {
    const double xStart = -0.496;
    const double yStart = yHoleCenter;
    const double zStart = zHoleCenter;
    const double tStart = 0.0;

    aval.AvalancheElectron(xStart, yStart, zStart, tStart);

    const auto& electrons = aval.GetElectrons();
    sumNeTot += electrons.size();

    if (!electrons.empty()) {
      const auto& primary = electrons.front();
      for (size_t k = 1; k < primary.path.size(); ++k) {
        if (primary.path[k - 1].x > xMeshPlane &&
            primary.path[k].x     <= xMeshPlane) {
          ++nPrimaryCrossedMesh;
          break;
        }
      }
    }

    int j = 0;
    for (const auto& el : electrons) {
      if (el.path.empty()) {
        ++j;
        continue;
      }

      const auto& start = el.path.front();
      const auto& end   = el.path.back();

      out << i << "  " << j++ << "  "
          << start.x << " " << start.y << " " << start.z << "   "
          << end.x   << " " << end.y   << " " << end.z   << "   "
          << el.status << "\n";

      const double x2 = end.x;
      const double y2 = end.y;
      const double z2 = end.z;

      if (x2 <= xAnodePlane + eps) {
        ++sumCollectedOnAnode;
      } else if (nearMeshX(x2) && !insideHoleYZ(y2, z2)) {
        ++sumLostOnMesh;
      } else if (x2 >= xCathPlane - eps) {
        ++sumLostOnCathode;
      } else {
        ++sumOtherEndpoints;
      }
    }
  }

  out.close();

  // --------------------------------------------------
  // 11. Summary
  // --------------------------------------------------
  const double Tmesh  = double(nPrimaryCrossedMesh) / nPrim;
  const double Gtotal = double(sumNeTot)            / nPrim;
  const double Ganode = double(sumCollectedOnAnode) / nPrim;

  std::cout << "\n=== Results (" << nPrim << " primary electrons) ===\n";
  std::cout << "Primary mesh transmission  Tmesh  = " << Tmesh  << "\n";
  std::cout << "Mean total avalanche size  Gtotal = " << Gtotal << "\n";
  std::cout << "Mean electrons on anode    Ganode = " << Ganode << "\n";
  std::cout << "Endpoints on anode                = " << sumCollectedOnAnode << "\n";
  std::cout << "Endpoints lost on mesh            = " << sumLostOnMesh       << "\n";
  std::cout << "Endpoints lost on cathode         = " << sumLostOnCathode    << "\n";
  std::cout << "Other endpoints                   = " << sumOtherEndpoints   << "\n";
  std::cout << "Saved endpoints to electron_endpoints.txt\n";

  return 0;
}