#include <cstdlib>
#include <iostream>
#include <fstream>

#include <TApplication.h>
#include <TCanvas.h>
#include <TH1F.h>

#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/TrackHeed.hh"
#include "Garfield/DriftLineRKF.hh"

#define TwoPi 6.18

using namespace Garfield;
using namespace std;

int main(int argc, char * argv[])
{
  TApplication app("app", &argc, argv);

    double Dim_x =  4;
    double Dim_y = 1.5;
    
    double Dim_z_down = 0;
    double Dim_z_up = 11.6;

  //plotfield
  bool plotField = true;
  bool plotDrift = true;
  bool plotContour = true;
  bool fieldDebug = false;

  //set gas
  const int ncoll = 10;
  MediumMagboltz gas;
//   gas.SetComposition("Xe");
// //   gas.SetComposition("Ar",90,"CH4",10);
//   gas.SetTemperature(293.15);
//   gas.SetPressure(760);

//   // Set the field range to be covered by the gas table. 
//   // const size_t nE = 11;
//   const size_t nE = 26;
//   const double emin = 0.;
//   const double emax = 5000.;
// //   Flag to request logarithmic spacing.
//   constexpr bool useLog = false;
//   gas.SetFieldGrid(emin, emax, nE, useLog);

//   gas.GenerateGasTable(ncoll);
// //   gas.WriteGasFile("ar_90_ch4_10.gas");
// //   gas.LoadGasFile("ar_90_ch4_10.gas");
// gas.WriteGasFile("xe.gas");
//   gas.LoadGasFile("xe.gas");
  gas.LoadGasFile("xe_origin.gas");
  gas.Initialise(true);
  auto installdir = std::getenv("GARFIELD_INSTALL");

  //read the ion mobility table from file
  if(!installdir)
  {std::cerr << "GARFIELD_INSTALL variable not set.\n";
  return 1;
  }
  const std::string path = installdir;
  gas.LoadIonMobility(path + "/share/Garfield/Data/IonMobility_Xe+_P12_Xe.txt");

  const string file_field = "/home/wxj/workspace/gar_practice/pin/field.txt";
  const string file_mesh = "/home/wxj/workspace/gar_practice/pin/mesh.mphtxt";
  const string file_die = "/home/wxj/workspace/gar_practice/pin/dielectrics.dat";
    
  ComponentComsol fm;
  fm.Initialise(file_mesh.c_str(), file_die.c_str(), file_field.c_str(), "mm");

  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();
      
	Sensor sensor;
	sensor.Clear();
	
	sensor.AddComponent(&fm);
    sensor.SetArea(-3,-3,-3,3,3,3);
              
	ViewField* fieldViewXY = new ViewField();
	ViewField* fieldViewXZ = new ViewField();
    ViewFEMesh* meshView = new ViewFEMesh();
	
if (plotField) 
{
	TCanvas *cFieldXY = new TCanvas("cFieldXY", "", 900, 800);
	fieldViewXY->SetComponent(&fm);
    fieldViewXY->SetPlaneXY();
	fieldViewXY->SetArea(-Dim_x, -Dim_y, Dim_x, Dim_y);
	fieldViewXY->SetSensor(&sensor);
	fieldViewXY->SetNumberOfSamples2d(500,500);
	fieldViewXY->SetCanvas(cFieldXY);
	fieldViewXY->Plot("e","colz");
    cFieldXY->Print("/home/wxj/workspace/gar_practice/pin/build/field_col.png");
    cFieldXY->Close();
}

if(plotContour)
{
  TCanvas *cFieldLine = new TCanvas("cFieldLine","cFieldLine",800,800);
  meshView->SetArea(-Dim_x, -Dim_y, Dim_x, Dim_y); 
  meshView->SetCanvas(cFieldLine);
  meshView->SetComponent(&fm);
  meshView->SetPlaneXY();
  meshView->SetFillMesh(true);
  meshView->SetColor(2, kGray);
  meshView->Plot(true);

  fieldViewXY->SetComponent(&fm);
  fieldViewXY->SetPlaneXY();
	fieldViewXY->SetArea(-Dim_x, -Dim_y, Dim_x, Dim_y);
	fieldViewXY->SetSensor(&sensor);
	fieldViewXY->SetNumberOfSamples2d(500,500);
	fieldViewXY->SetCanvas(cFieldLine);
	fieldViewXY->PlotContour();
  cFieldLine->Print("/home/wxj/workspace/gar_practice/pin/build/potential.png");
  cFieldLine->Close();
}

TrackSrim *tr = new TrackSrim();
tr->SetSensor(&sensor);
// Read SRIM output from file.
const std::string file = "/home/wxj/workspace/gar_practice/pin/Helium in Xenon (gas).txt";
tr->ReadFile(file);
// Read SRIM output from file
if (!tr->ReadFile(file)) {
    std::cerr << "Reading SRIM file failed.\n";
    return 0;
}

// Set the initial kinetic energy of the particle (in eV).
tr->SetKineticEnergy(5.486e6);//Am241
// Specify how many electrons we want to be grouped to a cluster.
tr->SetTargetClusterSize(1);
tr->SetClustersMaximum(20);
// tr->SetWorkFunction(228000);//放大10000倍
tr->SetFanoFactor(0.3);
tr->EnableLongitudinalStraggling(false);
tr->EnableTransverseStraggling(false);

// //电子轨迹
AvalancheMC* aval = new AvalancheMC();
aval->SetSensor(&sensor);

// Initial position
double x0 = -2., y0 = 0.8, z0 = 1., t0 = 0.;
// Direction of the track.
double dx0 = 1., dy0 = 0., dz0 = 0.;
// tr->NewTrack(x0, y0, z0, t0, dx0, dy0, dz0);

// ========== 关键修改：循环获取所有簇团 ==========
ViewDrift *driftView = new ViewDrift();
driftView->SetClusterMarkerSize(0.1);
driftView->SetCollisionMarkerSize(0.1);
driftView->SetColourElectrons(kRed);
driftView->SetColourTracks(kGreen + 3);
tr->EnablePlotting(driftView);
tr->EnableDebugging();
aval->EnablePlotting(driftView);

std::cout << "--------------- New Track ---------------" << std::endl;

// NewTrack返回bool，表示是否成功生成轨迹
bool trackStarted = tr->NewTrack(x0, y0, z0, t0, dx0, dy0, dz0);
if (!trackStarted) {
    std::cerr << "Failed to create new track.\n";
    return 1;
}

// 获取所有簇团
const auto& clusters = tr->GetClusters();
std::cout << "Number of clusters: " << clusters.size() << std::endl;
int totalElectrons = 0;//原初电离电子数量
int collectElectrons = 0;//被收集的电子的数量

TH1D *hx_end = new TH1D();//记录电子最后点的坐标
TH1D *hy_end = new TH1D();
TH1D *hz_end = new TH1D();

// 遍历所有簇团，对每个电子进行漂移模拟
for (size_t i = 0; i < clusters.size(); ++i) {
    const auto& cluster = clusters[i];
    std::cout << "Cluster " << i << ": position (" 
              << cluster.x << ", " << cluster.y << ", " << cluster.z 
              << "), electrons: " << cluster.n << std::endl;
    
    totalElectrons+=cluster.n;

// 对每个簇团中的电子进行漂移模拟
  for (int j = 0; j < cluster.n; ++j) 
  {
    // 电子漂移（注意：电子电荷为-1）
    double x0 = cluster.x;
    double y0 = cluster.y;
    double z0 = cluster.z;
    double t0 = 0.;
    double e0 = 0.1;
    cout <<"Electron started at ("<< x0 << ", " << y0 << ", " << z0 << ")\n";

    aval->DriftElectron(x0,y0,z0,t0);
    // 获取这个电子的终点
    double x1, y1, z1, e1, t1;
    int status;
    aval->GetElectronEndpoint(j, x0, y0, z0, t0, x1, y1, z1, t1, status);

    hx_end->Fill(x1);
    hy_end->Fill(y1);
    hz_end->Fill(z1);
    if((fabs(y1-0.1)>1e-6))
    {
      collectElectrons+=1;
      cout<<"y is "<<y1<<" ,collectElectrons are "<<collectElectrons<<endl;
    }
    std::cout << "Electron stopped at ("<< x1 << ", " << y1 << ", " << z1 << ")\n";
    std::cout << "Status code: " << status << "\n";
  }
}
double eff = ((double)collectElectrons/(double)totalElectrons)*100;
cout<<"Primaray Electron's number are "<<totalElectrons<<" , Collected Electron'number are "<<collectElectrons<<" , The Collected efficiency is "<< eff <<"%"<<endl;

if (plotDrift)
{
	TCanvas *cDrift = new TCanvas("cDrift", "", 600, 600);
	
  driftView->SetArea(-Dim_x, -Dim_y, Dim_x, Dim_y);
  driftView->SetPlaneXY();
	driftView->SetCanvas(cDrift);

  meshView->SetArea(-Dim_x, -Dim_y, Dim_x, Dim_y); 
  meshView->SetCanvas(cDrift);
  meshView->SetComponent(&fm);
  meshView->SetPlaneXY();
  meshView->SetFillMesh(true);
  meshView->SetColor(2, kGray);
  // meshView->Plot(true);
  driftView->Plot(true,true);
  meshView->Plot(true);
  cDrift->Update();
}
  
  app.Run();
}