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

#include <TApplication.h>
#include <TROOT.h>
#include <TCanvas.h>
#include <TFile.h>
#include <TTree.h>
#include <TSystem.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 <TLatex.h>


using namespace Garfield;
using namespace std;


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

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

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

  double Dim_x = 2.;
  double Dim_y = 2.;

  double Dim_z_dowm = -1.5;
  double Dim_z_up = 10; //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.;

  int Total_counts = 0;
  int Counts = 0;
  
  //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/zjj2536/TPC6PAD/build/outdata.txt");

  MediumMagboltz gas;
  gas.LoadGasFile("P10_1000.gas");
  //const std::string path = getenv("GARFIELD_HOME");
  gas.LoadIonMobility("/home/zjj2536/garfieldpp-master/Data/IonMobility_Ar+_Ar.txt");

  // Load the field map.
  ComponentComsol fm;
  fm.Initialise("mesh.mphtxt", "dielectrics.dat", "field.txt", "mm");
  fm.SetWeightingPotential("cathode.txt","c");
  fm.SetWeightingPotential("pad0.txt","p0");
  fm.SetWeightingPotential("pad1.txt","p1");
  fm.SetWeightingPotential("pad2.txt","p2");
  fm.SetWeightingPotential("pad3.txt","p3");
  fm.SetWeightingPotential("pad4.txt","p4");
  fm.SetWeightingPotential("pad5.txt","p5");
  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 = -500.;
  constexpr double tMax = 3500.; //ns
  constexpr double tStep = 2.;
  constexpr int nv = int(-tMin/tStep);
  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,"c");
  sensor.AddElectrode(&fm,"p0");
  sensor.AddElectrode(&fm,"p1");
  sensor.AddElectrode(&fm,"p2");
  sensor.AddElectrode(&fm,"p3");
  sensor.AddElectrode(&fm,"p4");
  sensor.AddElectrode(&fm,"p5");
  sensor.SetTimeWindow(tMin,tStep,nTimeBins);

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

  TCanvas* cFieldXZ = new TCanvas("cFieldXZ","",800,800);
  TCanvas* cFieldXY = new TCanvas("cFieldXY","",800,800);

  if(plotField){
      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");

      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");
      
      // cFieldWF = new TCanvas("cFieldWF","",800,800);
      // fieldViewWF->SetComponent(&fm);
      // fieldViewWF->SetPlane(0,-0.1,0,0,0.05,0);
      // fieldViewWF->SetArea(-Dim_x,Dim_z_dowm,Dim_x,Dim_z_up);
      // fieldViewWF->SetSensor(&sensor);
      // fieldViewWF->SetNumberOfSamples2d(500,500);
      // fieldViewWF->SetCanvas(cFieldWF);
      // fieldViewWF->PlotWeightingField(label_anodeT);

  }

  //track
  TrackSrim* tr = new TrackSrim();
  tr->SetSensor(&sensor);
  tr->ReadFile("P10alpha.txt");   
  tr->SetKineticEnergy(5.846e6); 
  tr->SetWorkFunction(27.0); //27eV*1000   
  tr->SetFanoFactor(0.19);

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


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

  aval->UseWeightingPotential(true);

  
  ViewDrift* driftView = new ViewDrift();
  //TCanvas* cDrift = new TCanvas("cDrift", "", 600, 600);

  if(plotDrift){
      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 = 50; 
  TLatex label;
  int count = 0;
  int ne = 0;

  for(unsigned int j = 0; j < nTrack; ++j){
      cout<<"----------New Track-----------"<<endl;
       rand = rrand.Uniform(-1,1);
       x0 = 0.5*rand;   
       y0 = 0.5*rand;

       z0 = 6.89;

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

       tr->NewTrack(x0,y0,z0,t0,dx,dy,-dz);

       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;
            ne = int(nc/300);

            //loop electrons
            for(int i = 0; i < nc; i++){
                
                //aval->SetElectronSignalScalingFactor(nc); //Multiplier of amplified signal
                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++;
                }
            }

       }
  }

  TCanvas* cDriftXZ = new TCanvas("cDriftXZ", "", 600, 600);
  TCanvas* cDriftXY = new TCanvas("cDriftXY", "", 600, 600);
  //TCanvas* cDriftYZ = new TCanvas("cDriftYZ", "", 600, 600);

  if(plotDrift){
      driftView->SetPlane(0,-0.1,0,0,0,0);
      driftView->SetArea(-Dim_x,Dim_z_dowm,Dim_x,Dim_z_up);
      driftView->SetCanvas(cDriftXZ);
      //driftView->SetColourElectrons(kRed);
      driftView->SetColourTracks(kGreen+3);
      driftView->Plot(twod,twod);

      driftView->SetPlane(0,0,-0.1,0,0,0);
      driftView->SetArea(-Dim_x,-Dim_y,Dim_x,Dim_y);
      driftView->SetCanvas(cDriftXY);
      //driftView->SetColourElectrons(kRed);
      driftView->SetColourTracks(kGreen+3);
      driftView->Plot(twod,twod);

      // driftView->SetPlane(0.1,0,0,0,0,0);
      // driftView->SetArea(-Dim_y,Dim_z_dowm,Dim_y,Dim_z_up);
      // driftView->SetCanvas(cDriftYZ);
      // driftView->SetColourElectrons(kRed);
      // driftView->SetColourTracks(kGreen+3);
      // driftView->Plot(twod,twod);
  }


  ofstream signal_pad0;
  signal_pad0.open("signal_pad0.txt");
  ofstream signal_pad1;
  signal_pad1.open("signal_pad1.txt");
  ofstream signal_pad2;
  signal_pad2.open("signal_pad2.txt");
  ofstream signal_pad3;
  signal_pad3.open("signal_pad3.txt");
  ofstream signal_pad4;
  signal_pad4.open("signal_pad4.txt");
  ofstream signal_pad5;
  signal_pad5.open("signal_pad5.txt");
  
  double q_pad0 = 0;
  double q_pad1 = 0;
  double q_pad2 = 0;
  double q_pad3 = 0;
  double q_pad4 = 0;
  double q_pad5 = 0;

  for(int i = 0; i < nTimeBins; i++){
 
      q_pad0 = (sensor.GetElectronSignal("p0",i));
      signal_pad0<<i<<"\t"<<q_pad0<<endl;

      q_pad1 = (sensor.GetElectronSignal("p1",i));
      signal_pad1<<i<<"\t"<<q_pad1<<endl;

      q_pad2 = (sensor.GetElectronSignal("p2",i));
      signal_pad2<<i<<"\t"<<q_pad2<<endl;

      q_pad3 = (sensor.GetElectronSignal("p3",i));
      signal_pad3<<i<<"\t"<<q_pad3<<endl;

      q_pad4 = (sensor.GetElectronSignal("p4",i));
      signal_pad4<<i<<"\t"<<q_pad4<<endl;

      q_pad5 = (sensor.GetElectronSignal("p5",i));
      signal_pad5<<i<<"\t"<<q_pad5<<endl;
     
  }

  ViewSignal Signalview_pad0;
  ViewSignal Signalview_pad1;
  ViewSignal Signalview_pad2;
  ViewSignal Signalview_pad3;
  ViewSignal Signalview_pad4;
  ViewSignal Signalview_pad5;

  TCanvas* ChargeSignal_pad0 = new TCanvas("pad0 signal", "", 600, 600);
  TCanvas* ChargeSignal_pad1 = new TCanvas("pad1 signal", "", 600, 600);
  TCanvas* ChargeSignal_pad2 = new TCanvas("pad2 signal", "", 600, 600);
  TCanvas* ChargeSignal_pad3 = new TCanvas("pad3 signal", "", 600, 600);
  TCanvas* ChargeSignal_pad4 = new TCanvas("pad4 signal", "", 600, 600);
  TCanvas* ChargeSignal_pad5 = new TCanvas("pad5 signal", "", 600, 600);

  if(plotSignal){
      
      Signalview_pad0.SetSensor(&sensor);
      sensor.IntegrateSignal("p0");
      Signalview_pad0.SetCanvas(ChargeSignal_pad0);
      Signalview_pad0.SetLabelY("Pad0 Induced Charge[fC]");
      Signalview_pad0.PlotSignal("p0","i");

      Signalview_pad1.SetSensor(&sensor);
      sensor.IntegrateSignal("p1");
      Signalview_pad1.SetCanvas(ChargeSignal_pad1);
      Signalview_pad1.SetLabelY("Pad1 Induced Charge[fC]");
      Signalview_pad1.PlotSignal("p1","i");

      Signalview_pad2.SetSensor(&sensor);
      sensor.IntegrateSignal("p2");
      Signalview_pad2.SetCanvas(ChargeSignal_pad2);
      Signalview_pad2.SetLabelY("Pad2 Induced Charge[fC]");
      Signalview_pad2.PlotSignal("p2","i");

      Signalview_pad3.SetSensor(&sensor);
      sensor.IntegrateSignal("p3");
      Signalview_pad3.SetCanvas(ChargeSignal_pad3);
      Signalview_pad3.SetLabelY("Pad3 Induced Charge[fC]");
      Signalview_pad3.PlotSignal("p3","i");

      Signalview_pad4.SetSensor(&sensor);
      sensor.IntegrateSignal("p4");
      Signalview_pad4.SetCanvas(ChargeSignal_pad4);
      Signalview_pad4.SetLabelY("Pad4 Induced Charge[fC]");
      Signalview_pad4.PlotSignal("p4","i");

      Signalview_pad5.SetSensor(&sensor);
      sensor.IntegrateSignal("p5");
      Signalview_pad5.SetCanvas(ChargeSignal_pad5);
      Signalview_pad5.SetLabelY("Pad5 Induced Charge[fC]");
      Signalview_pad5.PlotSignal("p5","i");

  }

   ViewField fieldView;
   TCanvas* cf = new TCanvas("cf", "", 600, 600);

  if (plotField) {
   
    fieldView.SetComponent(&fm);
    fieldView.SetPlane(0, -1, 0, 0, 0, 0);
    //fieldView.PlotProfile(-Dim_x, -Dim_y, Dim_z_dowm, Dim_x, Dim_y, Dim_z_up);
    cf->SetLeftMargin(0.16);
    fieldView.SetCanvas(cf);
    fieldView.PlotContour();
  }

  if(savefile){
    TFile* hist = new TFile("results6pad.root","RECREATE");
    ChargeSignal_pad0->Write("pad0_signal");
    ChargeSignal_pad1->Write("pad1_signal");
    ChargeSignal_pad2->Write("pad2_signal");
    ChargeSignal_pad3->Write("pad3_signal");
    ChargeSignal_pad4->Write("pad4_signal");
    ChargeSignal_pad5->Write("pad5_signal");
    cFieldXZ->Write("Field_XZ_Plane");
    cFieldXY->Write("Field_XY_Plane");
    cDriftXZ->Write("Drift_XZ");
    cDriftXY->Write("Drift_XY");
    //cDriftYZ->Write("Drift_YZ");
    cf->Write("EField_line");
    hist->Close();
  }

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

  app.Run(kTRUE);
  return 0;

}