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

#include <TCanvas.h>
#include <TROOT.h>
#include <TSystem.h>
#include <TApplication.h>
#include <TH1D.h>

#include "Garfield/MediumDiamond.hh"
#include "Garfield/ComponentConstant.hh"
#include "Garfield/ComponentUser.hh"
#include "Garfield/ComponentAnalyticField.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/TrackSrim.hh"
#include "Garfield/AvalancheMC.hh"

#include "Garfield/ViewDrift.hh"
#include "Garfield/ViewField.hh"
#include "Garfield/Plotting.hh"

#include "Garfield/FundamentalConstants.hh"
#include "Garfield/Random.hh"

using namespace Garfield;


bool readTransferFunction(Sensor& sensor) {

  std::ifstream infile;
  infile.open("ft.txt", std::ios::in);
  if (!infile) {
    std::cerr << "Could not read transfer function.\n";
    return false;
  }
  std::vector<double> times;
  std::vector<double> values;
  while (!infile.eof()) {
    double t = 0., f = 0.;
    infile >> t >> f;
    if (infile.eof() || infile.fail()) break;
    times.push_back(t);
    values.push_back(f);
  }
  infile.close();
  sensor.SetTransferFunction(times, values);
  return true;
}


int main(int argc, char * argv[]) {

  TApplication app("app", &argc, argv);

  // Define the medium.
  MediumDiamond si;
  si.SetTemperature(293.);

  // Make a plot of the drift velocities.
  plottingEngine.SetDefaultStyle();
  constexpr bool plotVelocity = true;
  if (plotVelocity) {
    si.PlotVelocity("eh", new TCanvas("cM", "", 600, 600));
  }
 
  // Sensor thickness [cm]
  constexpr double d = 70.e-4;

  // Make a component with constant drift field and weighting field.
  // Bias voltage [V]
  constexpr double vbias = -500.;
  ComponentConstant uniformField;
  uniformField.SetArea(-2 * d, 0., -2 * d, 2 * d, d, 2 * d);
  uniformField.SetMedium(&si);
  uniformField.SetElectricField(0, vbias / d, 0);
  uniformField.SetWeightingField(0, -1. / d, 0, "pad");

  // Depletion voltage [V]
  constexpr double vdep = -400.;
  // Make a component with linear drift field.
  auto eLinear = [d,vbias,vdep](const double /*x*/, const double y, 
                                const double /*z*/,
                                double& ex, double& ey, double& ez) {
    ex = ez = 0.;
    ey = (vbias - vdep) / d + 2 * y * vdep / (d * d);  
  };
  ComponentUser linearField;
  linearField.SetArea(-2 * d, 0., -2 * d, 2 * d, d, 2 * d);
  linearField.SetMedium(&si);
  linearField.SetElectricField(eLinear);
  // std::string efield = "ey = " + std::to_string((vbias - vdep) / d) + 
  //                      " + 2 * y * " + std::to_string(vdep / (d * d));
  // linearField.SetElectricField(efield);

  // Make a component with analytic weighting field for a strip or pixel.
  constexpr double pitch = 55.e-4;
  constexpr double halfpitch = 0.5 * pitch;
  ComponentAnalyticField wField;
  wField.SetMedium(&si);
  wField.AddPlaneY(0, vbias, "back");
  wField.AddPlaneY(d, 0, "front");
  wField.AddStripOnPlaneY('z', d, -halfpitch, halfpitch, "strip");
  wField.AddPixelOnPlaneY(d, -halfpitch, halfpitch, 
                             -halfpitch, halfpitch, "pixel");

  // Create a sensor. 
  Sensor sensor;
  sensor.AddComponent(&linearField); 
  const std::string label = "strip";
  sensor.AddElectrode(&wField, label);

  // Plot the drift field if requested.
  constexpr bool plotField = true;
  if (plotField) {
    ViewField* fieldView = new ViewField(); 
    fieldView->SetSensor(&sensor); 
    fieldView->SetArea(-0.5 * d, 0, 0.5 * d, d);
    fieldView->PlotContour("ey");
  }
  // Plot the weighting potential if requested.
  constexpr bool plotWeightingField = true;
  if (plotWeightingField) {
    ViewField* wfieldView = new ViewField(); 
    wfieldView->SetComponent(&wField); 
    wfieldView->SetArea(-0.5 * d, 0, 0.5 * d, d);
    wfieldView->PlotContourWeightingField("strip", "v");
  }

  Shaper shaper(1, 25., 1., "bipolar");
  sensor.SetTransferFunction(shaper);
 /*   auto fT = [](const double t) {
      constexpr double tau = 25.;
      return (t / tau) * exp(1 - t / tau);
    };
   sensor.SetTransferFunction(fT);
*/
  // Set the time bins.
  const unsigned int nTimeBins = 1000;
  const double tmin =  0.;
  const double tmax = 10.;
  const double tstep = (tmax - tmin) / nTimeBins;
  sensor.SetTimeWindow(tmin, tstep, nTimeBins);
  //if (!readTransferFunction(sensor)) return 0;
  //sensor.ClearSignal();


 /* // Set up Heed.
  TrackHeed track(&sensor);
  // Set the particle type and momentum [eV/c].
  track.SetParticle("pion");
  track.SetMomentum(180.e9);
*/
  // Simulate electron/hole drift lines using MC integration.
  // Create a track class and connect it to a sensor.
  TrackSrim track(&sensor);
   // Read SRIM output from file.
    const std::string file = "aindiamond.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(3.e6);
    const double za = 6. / 12.;
    track.SetAtomicMassNumbers(6. / za, 6);
    track.SetDensity(3.52);
    // Specify how many electrons we want to be grouped to a cluster.
    track.SetTargetClusterSize(900);
    // tr.SetClustersMaximum(1000);

    // Make some plots of the SRIM data.
    track.PlotEnergyLoss();
    track.PlotRange();
    track.PlotStraggling();
    // Print a table of the SRIM data.
    track.Print();
    
    AvalancheMC drift(&sensor);
    // Use steps of 1 micron.
    drift.SetDistanceSteps(1.e-4);
    
  // Plot the signal if requested.
  constexpr bool plotSignal = true;
  TCanvas* cSignal = nullptr;
  if (plotSignal) { 
    cSignal = new TCanvas("cSignal", "", 600, 600);
  }

  constexpr bool plotDrift = true;
  ViewDrift* driftView = nullptr;
  TCanvas* cDrift = nullptr;
  if (plotDrift) {
    cDrift = new TCanvas("cDrift", "", 600, 600);
    driftView = new ViewDrift();
    driftView->SetArea(-0.5 * d, 0, -0.5 * d, 0.5 * d, d, 0.5 * d);
    driftView->SetCanvas(cDrift);
    track.EnablePlotting(driftView);
  }
  // Flag to randomise the position of the track.  
  constexpr bool smearx = true; 
  constexpr unsigned int nEvents = 3;
  // Flag to save the signal to a file.
  constexpr bool writeSignal = true;
    for (unsigned int i = 0; i < nEvents; ++i) {
        if (plotDrift) driftView->Clear();
        // Reset the signal.
        sensor.ClearSignal();
        if (i % 10 == 0) std::cout << i << "/" << nEvents << "\n";
        // Simulate a charged-particle track.
        double xt = 0.;
        if (smearx) xt = -0.5 * pitch + RndmUniform() * pitch;
        const unsigned int nTracks = 10;
        for (unsigned int i = 0; i < nTracks; ++i) {
            if (!track.NewTrack(xt, d, 0., 0., 0., 1., 0.)) {
                std::cerr << "Generating clusters failed; skipping this track.\n";
                continue;
            }
            // Retrieve the clusters.
            const auto& clusters = track.GetClusters();
            if (clusters.empty()) continue;
            // Count the total number of electrons.
            unsigned int netot = 0;
            for (const auto& cluster : clusters) netot += cluster.n;
            const auto& last = clusters.back();
        }
        
        
        
        /* track.NewTrack(xt, 0, 0, 0, 0, 1, 0);
         // Retrieve the clusters along the track.
         for (const auto& cluster : track.GetClusters()) {
         // Loop over the electrons in the cluster.
         for (const auto& electron : cluster.n) {
         // Simulate the electron and hole drift lines.
         if (plotDrift) {
         drift.DisablePlotting();
         if (RndmUniform() < 0.01) drift.EnablePlotting(driftView);
         }
         drift.DriftElectron(electron.x, electron.y, electron.z, electron.t);
         drift.DriftHole(electron.x, electron.y, electron.z, electron.t);
         }
         }
         
         
         const unsigned int nTracks = 5;
         for (unsigned int i = 0; i < nTracks; ++i) {
         if (!track.NewTrack(0., 0., 0., 0., 1., 0., 0.)) {
         std::cerr << "Generating clusters failed; skipping this track.\n";
         continue;
         }
         // Retrieve the clusters.
         const auto& clusters = track.GetClusters();
         if (clusters.empty()) continue;
         // Count the total number of electrons.
         unsigned int netot = 0;
         for (const auto& cluster : clusters) netot += cluster.n;
         const auto& last = clusters.back();
         }
         */
        sensor.ConvoluteSignals();
        if (plotSignal) {
            sensor.PlotSignal(label, cSignal);
            cSignal->Update();
            gSystem->ProcessEvents();
        }
        
        if (plotDrift) {
            constexpr bool twod = true;
            driftView->Plot(twod);
            cDrift->Update();
            gSystem->ProcessEvents();
        }
        
        // Save the induced current signal to a file.
        // constexpr bool writeSignal = true;
        if (writeSignal) {
            char filename[50];
            sprintf(filename, "signal_%05d.txt", i);
            std::ofstream outfile;
            outfile.open(filename, std::ios::out);
            for (unsigned int j = 0; j < nTimeBins; ++j) {
                const double t = (j + 0.5) * tstep;
                const double f = sensor.GetSignal(label, j);
                const double fe = sensor.GetElectronSignal(label, j);
                const double fh = sensor.GetIonSignal(label, j);
                outfile << t << "  " << f << "  " << fe << "  " << fh << "\n";
            }
            outfile.close();
        }
    }

  if (plotVelocity || plotSignal ||
      //plotDrift ||
      plotField || plotWeightingField) {
    app.Run();
  }
}