#include <iostream>
#include <fstream>
#include <cmath>
#include "math.h"

#include <TApplication.h>
#include <TROOT.h>
#include <TCanvas.h>
//#include <Tfile.h>

#include "Garfield/MediumConductor.hh"
#include "Garfield/ViewGeometry.hh"
#include "Garfield/SolidBox.hh"
#include "Garfield/ComponentAnsys123.hh"
#include "Garfield/ComponentComsol.hh"
#include "Garfield/ViewField.hh"
#include "Garfield/ViewFEMesh.hh"
#include "Garfield/MediumMagboltz.hh"
#include "Garfield/Sensor.hh"
#include "Garfield/AvalancheMicroscopic.hh"
#include "Garfield/AvalancheMC.hh"
#include "Garfield/Random.hh"
#include "Garfield/TrackSrim.hh"
#include "Garfield/GeometrySimple.hh"
#include "Garfield/ViewSignal.hh"
#include "Garfield/ViewCell.hh"
#include "Garfield/Plotting.hh"
#include <TSystem.h>
#include <cstdlib>
#include <TLatex.h>
#include "Garfield/TrackHeed.hh"


using namespace Garfield;
using namespace std;


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

  TApplication app("app", &argc, argv);
  
  string filename;
  string mypath;
  string commend;


  constexpr bool plotField = true;
  constexpr bool plotDrift = true;
  constexpr bool plotSignal = true;
  constexpr bool savefile = true;

  double Dim_x = 20.;
  double Dim_y = 20.;

  double Dim_z_dowm = -25.;
  double Dim_z_up = 25.; //cm

  double x0 = 0.5;
  double y0 = 0.5;
  double z0 = 0.5;
  double t0 = 0.;
  double dx = 1.;
  double dy = 1.;
  double dz = 1.;

  double Total_counts;
  double Counts;
  
  //cluster
  double xc = 0., yc = 0., zc = 0., tc = 0., ec = 0., ekin = 0.;
  int nc = 0, status = 0;
   
  //Electronic start and end position
  double xe1 = 0., ye1 = 0., ze1 = 0., te1 = 0., ee1 = 0.;
  double xe2 = 0., ye2 = 0., ze2 = 0., te2 = 0., ee2 = 0.;

  ofstream outdata("/home/jack/Desktop/2022.5taosimulation/Srim/build/outdata.txt");


  TCanvas* ChargeSignal_Cathode = nullptr;
  TCanvas* ChargeSignal_anodeT = nullptr;
  TCanvas* ChargeSignal_anodeB = nullptr;
  TCanvas* cFieldXY = nullptr;
  TCanvas* cFieldXZ = nullptr;
  TCanvas* cDrift = nullptr;
  
  Total_counts = 0;
  Counts = 0;

  MediumMagboltz gas;
  gas.LoadGasFile("p160kpa.gas");
  const std::string path = getenv("GARFIELD_HOME");
  gas.LoadIonMobility(path + "/Data/IonMobility_Ar+_Ar.txt");

  const std::string label_cathode = "cathode";
  const std::string label_anodeT = "anodeT";
  const std::string label_anodeB = "anodeB";

    // Load the field map.
  ComponentComsol fm;
  fm.Initialise("mesh.mphtxt", "dielectrics.dat", "field.txt", "cm");
  fm.SetWeightingField("cathode.txt",label_cathode);
  fm.SetWeightingField("anodeT.txt",label_anodeT);
  fm.SetWeightingField("anodeB.txt",label_anodeB);
  fm.PrintRange();
   
    // Associate the gas with the corresponding field map material. 
  const unsigned int nMaterials = fm.GetNumberOfMaterials();
  for (unsigned int i = 0; i < nMaterials; ++i) {
    const double eps = fm.GetPermittivity(i);
    if (eps == 1.) fm.SetMedium(i, &gas);
  }
  fm.PrintMaterials();

  //signal time window
  constexpr double tMin = 0.;
  constexpr double tMax = 10000.; //ns
  constexpr double tStep = 1.;
  constexpr int nTimeBins = int((tMax-tMin)/tStep);
  
  Sensor sensor;
  sensor.Clear();

  sensor.AddComponent(&fm);
  sensor.SetArea(-Dim_x,-Dim_y,Dim_z_dowm,Dim_x,Dim_y,Dim_z_up);
  sensor.AddElectrode(&fm,label_cathode);
  sensor.AddElectrode(&fm,label_anodeT);
  sensor.AddElectrode(&fm,label_anodeB);
  sensor.SetTimeWindow(0.,tStep,nTimeBins);

  ViewField* fieldViewXY = new ViewField();
  ViewField* fieldViewXZ = new ViewField();

  if(plotField){
      cFieldXZ = new TCanvas("cFieldXZ","",800,800);
      fieldViewXZ->SetComponent(&fm);
      fieldViewXZ->SetPlane(0,-0.1,0,0,0.05,0);
      fieldViewXZ->SetArea(-Dim_x,Dim_z_dowm,Dim_x,Dim_z_up);
      fieldViewXZ->SetSensor(&sensor);
      fieldViewXZ->SetNumberOfSamples2d(500,500);
      fieldViewXZ->SetCanvas(cFieldXZ);
      fieldViewXZ->Plot("e","colz");

      cFieldXY = new TCanvas("cFieldXY","",800,800);
      fieldViewXY->SetComponent(&fm);
      fieldViewXY->SetPlane(0,0,-0.1,0,0,0.05);
      fieldViewXY->SetArea(-Dim_x,-Dim_y,Dim_x,Dim_y);
      fieldViewXY->SetSensor(&sensor);
      fieldViewXY->SetNumberOfSamples2d(500,500);
      fieldViewXY->SetCanvas(cFieldXY);
      fieldViewXY->Plot("e","colz");


  }

  //track
  TrackSrim* tr = new TrackSrim();
  tr->SetSensor(&sensor);
  tr->ReadFile("p160kpa.txt");   
  tr->SetKineticEnergy(5.846e6); 
  tr->SetWorkFunction(27000);    
  tr->SetFanoFactor(0.19);

  constexpr double za = 0.9*(18./40.) + 0.1*(10./16.);
  tr->SetAtomicMassNumbers(10./za, 10.);  

  tr->SetClustersMaximum(100); 
  tr->Print();


  //electron drift
  AvalancheMC* aval = new AvalancheMC();
  aval->SetSensor(&sensor);
  aval->EnableSignalCalculation(true);
  aval->UseWeightingPotential(true);
  aval->DisableAvalancheSizeLimit();

  
  ViewDrift* driftView = new ViewDrift();
  if(plotDrift){
      cDrift = new TCanvas("cDrift", "", 600, 600);

      driftView->SetClusterMarkerSize(0.1);
      driftView->SetCollisionMarkerSize(0.1);
      driftView->SetColourElectrons(kRed);
      driftView->SetColourTracks(kGreen + 3);
      tr->EnablePlotting(driftView);
      aval->EnablePlotting(driftView);
  }



  TRandom rrand;
  double rand;
  const unsigned int nTrack = 10; 
  TLatex label;
  int count = 0;

  for(unsigned int j = 0; j < nTrack; ++j){
      cout<<"----------New Track-----------"<<endl;
       rand = rrand.Uniform(0,360);
       x0 = 0.1*sin(rand*Pi);   
       y0 = 0.1*cos(rand*Pi);

       z0 = -4.74;

       dx = rrand.Uniform(-1,1); 
       dy = rrand.Uniform(-1,1);
       dz = rrand.Uniform(0,1);

       tr->NewTrack(0,0,z0,0,0,0,1);

       while(tr->GetCluster(xc,yc,zc,tc,nc,ec,ekin)){
           ee1 = 0.1; 
            cout<<"========================================="<<endl;
            cout<<"Number of Electron: "<<nc<<endl;
            cout<<"Position of Clusters: "<<xc<<","<<yc<<","<<zc<<endl;
            outdata<<"cluster position: "<<nc<<" "<<xc<<" "<<yc<<" "<<zc<<endl;

            //loop electrons
            for(int i = 0; i < nc; i++){
                // int ne = 0, ni = 0;
                // aval->GetAvalancheSize(ne,ni);
                
                aval->DriftElectron(xc,yc,zc,tc);   
                const int np = aval->GetNumberOfElectronEndpoints(); 
                cout<<"Multiplicity: "<<np<<endl;
                outdata<<"Multiplicity: "<<np<<endl;

                for(int j = 0; j < np; j++){
                    aval->GetElectronEndpoint(j,xe1,ye1,ze1,te1,xe2,ye2,ze2,te2,status); 
                    cout<<" Electron start position: "<<xe1<<","<<ye1<<","<<ze1<<"\t"<<"Electron stop position: "<<xe2<<","<<ye2<<","<<ze2<<endl;
                    outdata<<" Electron start position: "<<xe1<<","<<ye1<<","<<ze1<<"\t"<<"Electron stop position: "<<xe2<<","<<ye2<<","<<ze2<<endl;
                    Total_counts++;
                }
            }

       }
  }

  if(plotDrift){
      constexpr bool plotMesh = false;
      if(plotMesh){
          ViewFEMesh* meshView = new ViewFEMesh();
          meshView->SetArea(-Dim_x/20, -Dim_y/20, Dim_z_dowm, Dim_x/20, Dim_y/20, Dim_z_up);
          meshView->SetCanvas(cDrift);
          meshView->SetComponent(&fm);
          //x-z projection
          meshView->SetPlane(0,-0.1,0,0,0.05,0);
          meshView->SetFillMesh(true);
          driftView->SetColourElectrons(kRed);
          driftView->SetColourTracks(kGreen+3);
          //set color
          meshView->SetViewDrift(driftView);
          meshView->Plot();
      }else{
          driftView->SetPlane(0,-0.1,0,0,0,0);
          driftView->SetArea(-Dim_x/4.,Dim_z_dowm/4.,Dim_x/4.,Dim_z_up/4.);
          driftView->SetCanvas(cDrift);
          constexpr bool twod = true;
          driftView->Plot(twod,twod);
      }
  }

  ofstream signal_cathode;
  signal_cathode.open("signal_cathode.txt");

  ofstream signal_anodeT;
  signal_anodeT.open("signal_anodeT.txt");

  ofstream signal_anodeB;
  signal_anodeB.open("signal_anodeB.txt");

  double q_cathode = 0;
  double q_anodeT = 0;
  double q_anodeB = 0;

  for(int i = 0; i < nTimeBins; i++){
      q_cathode = (sensor.GetSignal(label_cathode,i));
      signal_cathode<<i<<"\t"<<q_cathode<<endl;

      q_anodeT = (sensor.GetSignal(label_anodeT,i));
      signal_anodeT<<i<<"\t"<<q_anodeT<<endl;

      q_anodeB = (sensor.GetSignal(label_anodeB,i));
      signal_anodeB<<i<<"\t"<<q_anodeB<<endl;

  }

  ViewSignal Signalview_cathode;
  ViewSignal Signalview_anodeT;
  ViewSignal Signalview_anodeB;

  if(plotSignal){
      ChargeSignal_Cathode = new TCanvas("cathode signal", "", 600, 600);
      Signalview_cathode.SetSensor(&sensor);
      sensor.IntegrateSignal(label_cathode);
      Signalview_cathode.SetCanvas(ChargeSignal_Cathode);
      Signalview_cathode.SetLabelY("Induced Charge[fC]");
      Signalview_cathode.PlotSignal(label_cathode);

      ChargeSignal_anodeT = new TCanvas("anodeT signal", "", 600, 600);
      Signalview_anodeT.SetSensor(&sensor);
      sensor.IntegrateSignal(label_anodeT);
      Signalview_anodeT.SetCanvas(ChargeSignal_anodeT);
      Signalview_anodeT.SetLabelY("Induced Charge[fC]");
      Signalview_anodeT.PlotSignal(label_anodeT,true,false,false,true,false);

      ChargeSignal_anodeB = new TCanvas("anodeB signal", "", 600, 600);
      Signalview_anodeB.SetSensor(&sensor);
      sensor.IntegrateSignal(label_anodeB);
      Signalview_anodeB.SetCanvas(ChargeSignal_anodeB);
      Signalview_anodeB.SetLabelY("Induced Charge[fC]");
      Signalview_anodeB.PlotSignal(label_anodeB,true,false,false,true,false);
  }

   ViewField fieldView;

  if (plotField) {
    fieldView.SetComponent(&fm);
    fieldView.SetPlane(0, -1, 0, 0, 0, 0);
    TCanvas* cf = new TCanvas("cf", "", 600, 600);
    cf->SetLeftMargin(0.16);
    fieldView.SetCanvas(cf);
    fieldView.PlotContour();
  }

  cout<<"================"<<endl;
  cout<<"Done"<<endl;

  app.Run(kTRUE);
  return 0;

}