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

#include <TROOT.h>
#include <TCanvas.h>
#include <TGraph.h>
#include <TMultiGraph.h>
#include <TLegend.h>
#include <TAxis.h>
#include <TStyle.h>

#include "Garfield/MediumMagboltz.hh"
#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/DriftLineRKF.hh"

using namespace Garfield;

int main() {
  gROOT->SetBatch(true);
  gStyle->SetOptStat(0);

  // ------------------------------------------------------------
  // Geometry from the article:
  // Ar 50% - CO2 50%
  // rc = 1.25 cm
  // ra = 24 micrometer
  // ------------------------------------------------------------
  const double rCathode = 1.25;        // cm
  const double rWire = 24.e-4;         // cm
  const double dWire = 2.0 * rWire;    // cm
  const double vCathode = 0.0;         // V

  const double temperature = 293.15;   // K

  // Four pressures from the article.
  std::vector<double> pressuresAtm = {
    0.40,
    0.79,
    1.19,
    1.78
  };

  // ROOT styles for curves.
  std::vector<int> colors = {
    kMagenta + 1,
    kBlue + 1,
    kRed + 1,
    kGreen + 2
  };

  std::vector<int> markers = {
    20,
    21,
    22,
    23
  };

  // ------------------------------------------------------------
  // Output CSV file
  // ------------------------------------------------------------
  std::ofstream out("gain_vs_voltage_4pressures.csv");
  out << "# Ar/CO2 = 50/50\n";
  out << "# rCathode = " << rCathode << " cm\n";
  out << "# rWire = " << rWire << " cm\n";
  out << "pressure_atm,V_anode_kV,mean_gain,rms_gain,n_ok\n";

  // ------------------------------------------------------------
  // Canvas and multigraph
  // ------------------------------------------------------------
  TCanvas canvas("canvas", "Gas gain vs anode voltage", 900, 700);
  canvas.SetGrid();
  canvas.SetLogy();

  TMultiGraph mg;
  TLegend legend(0.62, 0.18, 0.88, 0.40);
  legend.SetBorderSize(0);
  legend.SetFillStyle(0);

  // Store graphs so they do not disappear from memory.
  std::vector<TGraph*> graphs;

  // ------------------------------------------------------------
  // Starting points for averaging
  // ------------------------------------------------------------
  const int nStarts = 8;
  const double rStart = 0.90 * rCathode;
  const double pi = 3.14159265358979323846;

  // ------------------------------------------------------------
  // Loop over pressures
  // ------------------------------------------------------------
  for (size_t ip = 0; ip < pressuresAtm.size(); ++ip) {
    const double pressureAtm = pressuresAtm[ip];
    const double pressureTorr = pressureAtm * 760.0;

    std::cout << "\n========================================\n";
    std::cout << "Pressure = " << pressureAtm << " atm\n";
    std::cout << "========================================\n";

    // ------------------------------------------------------------
    // Gas for this pressure
    // ------------------------------------------------------------
    MediumMagboltz gas("ar", 50., "co2", 50.);
    gas.SetTemperature(temperature);
    gas.SetPressure(pressureTorr);

    // Useful for high electric fields near the wire.
    gas.SetMaxElectronEnergy(200.);

    // Penning transfer for Ar/CO2.
    gas.EnablePenningTransfer();

    // Initialise Magboltz.
    gas.Initialise(true);

    // ------------------------------------------------------------
    // Gas table for this pressure
    // ------------------------------------------------------------
    const double eMin = 10.;        // V/cm
    const double eMax = 300.e3;     // V/cm
    const int nE = 20;
    const bool useLog = true;

    gas.SetFieldGrid(eMin, eMax, nE, useLog);


    const int nColl = 10;
    gas.GenerateGasTable(nColl);

    // Optional gas table output for each pressure.
    std::string gasFileName =
      "ar50_co2_50_p" + std::to_string(int(pressureAtm * 100)) + "atm.gas";
    gas.WriteGasFile(gasFileName);

    std::vector<double> voltageKV;
    std::vector<double> meanGains;

    // ------------------------------------------------------------
    // Loop over anode voltage
    // ------------------------------------------------------------
    for (double vAnodeKV = 0.5; vAnodeKV <= 3.6 + 1.e-9; vAnodeKV += 0.1) {
      const double vAnode = vAnodeKV * 1000.0; // V

      ComponentAnalyticField cmp;
      cmp.SetMedium(&gas);

      // Central anode wire.
      cmp.AddWire(0., 0., dWire, vAnode, "s");

      // Cylindrical cathode.
      cmp.AddTube(rCathode, vCathode, 0);

      Sensor sensor;
      sensor.AddComponent(&cmp);

      DriftLineRKF drift(&sensor);

      // Important: enable multiplication.
      drift.EnableAvalanche(true);

      double sumGain = 0.;
      double sumGain2 = 0.;
      int nOk = 0;

      for (int i = 0; i < nStarts; ++i) {
        const double phi = 2.0 * pi * double(i) / double(nStarts);

        const double x0 = rStart * std::cos(phi);
        const double y0 = rStart * std::sin(phi);
        const double z0 = 0.;
        const double t0 = 0.;

        drift.DriftElectron(x0, y0, z0, t0);

        const double gain = drift.GetGain();

        if (std::isfinite(gain) && gain > 0.) {
          sumGain += gain;
          sumGain2 += gain * gain;
          ++nOk;
        }
      }

      double meanGain = 0.;
      double rmsGain = 0.;

      if (nOk > 0) {
        meanGain = sumGain / nOk;
        const double variance = sumGain2 / nOk - meanGain * meanGain;
        rmsGain = std::sqrt(std::max(0.0, variance));
      }

      std::cout << std::fixed << std::setprecision(2)
                << "p = " << pressureAtm << " atm"
                << "   V = " << vAnodeKV << " kV"
                << "   gain = " << std::scientific << meanGain
                << "   nOk = " << std::defaultfloat << nOk << "\n";

      out << std::fixed << std::setprecision(3)
          << pressureAtm << ","
          << vAnodeKV << ","
          << std::scientific << meanGain << ","
          << rmsGain << ","
          << nOk << "\n";

      if (meanGain > 0. && std::isfinite(meanGain)) {
        voltageKV.push_back(vAnodeKV);
        meanGains.push_back(meanGain);
      }
    }

    // ------------------------------------------------------------
    // Create graph for this pressure
    // ------------------------------------------------------------
    TGraph* gr = new TGraph(voltageKV.size(), voltageKV.data(), meanGains.data());

    gr->SetLineColor(colors[ip]);
    gr->SetMarkerColor(colors[ip]);
    gr->SetLineWidth(2);
    gr->SetMarkerStyle(markers[ip]);
    gr->SetMarkerSize(1.0);

    mg.Add(gr, "LP");

    std::string label = std::to_string(pressureAtm) + " atm";
    legend.AddEntry(gr, label.c_str(), "lp");

    graphs.push_back(gr);
  }

  out.close();

  // ------------------------------------------------------------
  // Draw final plot
  // ------------------------------------------------------------
  mg.SetTitle("Gas gain vs anode voltage;Anode voltage U_{a} [kV];Mean gas gain");
  mg.Draw("A");

  mg.GetXaxis()->CenterTitle();
  mg.GetYaxis()->CenterTitle();

  mg.GetXaxis()->SetLimits(0.4, 3.7);
  mg.SetMinimum(1.);
  mg.SetMaximum(2.e7);

  legend.Draw();

  canvas.SaveAs("gain_vs_voltage_4pressures.png");
  canvas.SaveAs("gain_vs_voltage_4pressures.pdf");

  std::cout << "\nSaved files:\n";
  std::cout << "  gain_vs_voltage_4pressures.csv\n";
  std::cout << "  gain_vs_voltage_4pressures.png\n";
  std::cout << "  gain_vs_voltage_4pressures.pdf\n";

  return 0;
}