#include "TRandom.h"      // ROOT random number generators, gRandom
#include "TFile.h"        // ROOT file I/O
#include "TTree.h"        // ROOT tree storage
#include "TH1.h"          // 1D histograms
#include "TH2.h"          // 2D histograms
#include "TStyle.h"       // ROOT plotting style controls
#include "TCanvas.h"      // ROOT canvas drawing
#include "TBenchmark.h"   // timing measurement
#include "TGraph.h"       // for energy loss interpolation
#include <cstring>
#include "TApplication.h" // ROOT app loop
#include "ClassTransfer.h" // Reaction kinematics and MC event generation
#include "ClassAnasen.h"   // ANASEN detector model classes (SX3, PW, etc.)
#include <stdio.h>
#include <stdlib.h>
#include <set>
#include "TLegend.h"
#include "TH1D.h"
#include "TObjArray.h"
#include "TBranch.h"
#include <iostream>
#include <fstream>

//======== Generate light particle based on reaction
// calculate real and reconstructed tracks and Q-value uncertainty

// Function to load energy loss table from file
TGraph* LoadELoss(const char* filename) {
  TGraph* g = new TGraph(filename, "%lg %lg");
  return g;
}

bool IsDeadAnode(int id){
    static std::set<int> dead = {}; // add dead anode IDs here, 0-23
    return dead.count(id);
}

bool IsDeadCathode(int id){
    static std::set<int> dead = {}; // add dead cathode IDs here, 0-23
    return dead.count(id);
}

bool IsDeadSX3(int id){ 
    static std::set<int> dead = {}; // add dead SX3 IDs here, 0-23 1,7,9,3
    return dead.count(id);
}

static std::set<pair<int,int>> ReactionProductb = { {1,1} }; // add reaction product b (light particle) A,Z pairs here, e.g. {1,1} for proton, {4,2} for alpha 

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

  printf("=========================================\n");
  printf("===       ANASEN Monte Carlo          ===\n");
  printf("=========================================\n");

  // number of events can be overridden from command line
  int numEvent = 1000000;
  if( argc >= 2 ) numEvent = atoi(argv[1]);

  // load energy loss tables (assume units: E in MeV, dE/dx in MeV/(mg/cm^2), density in mg/cm^3)
  TGraph* elossAlpha = LoadELoss("../ELoss/E_vs_x_alpha.dat"); // for light particle (alpha)
  TGraph* elossProton = LoadELoss("../ELoss/E_vs_x_proton.dat");     // for heavy particle (proton)
  TGraph *invgAlpha = new TGraph(elossAlpha->GetN(), elossAlpha->GetY(), elossAlpha->GetX());
  TGraph *invgProton = new TGraph(elossProton->GetN(), elossProton->GetY(), elossProton->GetX());

  
  //Plot energy loss tables (sanity check), vis will not work if this is ran without X11 display (e.g. on cluster), so comment out if running in headless mode
  auto c1 = new TCanvas("c1", "Graph Example", 800, 600);
  auto g = elossAlpha;
  g->SetTitle("Energy Loss Table (Alpha);cm;Kinetic Energy (MeV)");
  g->Draw("ALP");
  g->SetLineColor(kRed);
  //c1->SetLogy();
  //c1->SetLogx();
  c1->Print("eloss_alpha.png");

  auto c2 = new TCanvas("c2", "Graph Example", 800, 600);
  auto g2 = elossProton;
  g2->SetTitle("Energy Loss Table (Proton);cm;Kinetic Energy (MeV)");
  g2->Draw("ALP");
  g2->SetLineColor(kBlue);
  c2->Print("eloss_proton.png");

  // Reaction setup: projectile + target configuration, energy, and product IDs
  TransferReaction transfer;

  transfer.SetA(18, 10, 0);           // e.g., 24Mg (Z=12) with 0 excitation
  transfer.SetIncidentEnergyAngle((4.4), 0, 0); // arguments are KEA in MeV/u, theta and phi in degree
  transfer.Seta( 4, 2);              // identify reaction product a in internal indexing e.g., 4He (alpha)
  transfer.Setb(ReactionProductb.begin()->first, ReactionProductb.begin()->second);              // identify reaction product b e.g., 1H (proton)
  transfer.SetB(21, 11);              // identify reaction product B e.g., 23Na (Z=11)

  // TODO add alpha source or alternative reaction channel selection

  // Excited state lists (target and projectile/excited products)
  std::vector<float> ExAList = {0}; // projectile excitation states in MeV
  std::vector<float> ExList = {0}; // target excitation states in MeV

  // define vertex position uniform distribution ranges (mm)
  double vertexXRange[2] = {   -5,    5}; // mm
  double vertexYRange[2] = {   -5,    5};
  double vertexZRange[2] = { -100,  100};

  // detector resolution / uncertainty parameters
  double sigmaSX3_W = -1; // mm, if < 0 use mid-point (no spread in SX3 horizontal dimension)
  double sigmaSX3_L = 3;  // mm, vertical spread for SX3
  double sigmaPW_A  = 0;  // normalized anode uncertainty term (0-1)
  double sigmaPW_C  = 0;  // normalized cathode uncertainty term (0-1)

  // status printout
  printf("------------ Vertex :\n");
  printf("X : %7.2f - %7.2f mm\n", vertexXRange[0], vertexXRange[1]);
  printf("Y : %7.2f - %7.2f mm\n", vertexYRange[0], vertexYRange[1]);
  printf("Z : %7.2f - %7.2f mm\n", vertexZRange[0], vertexZRange[1]);
  printf("------------ Uncertainty :\n");
  printf(" SX3 horizontal : %.1f\n", sigmaSX3_W);
  printf(" SX3   vertical : %.1f\n", sigmaSX3_L);
  printf("          Anode : %.1f mm\n", sigmaPW_A);
  printf("        Cathode : %.1f mm\n", sigmaPW_C);
  printf("        num_eve : %d \n",numEvent);

  // calculates energy/momentum/kinematics constants for transfer reaction
  transfer.CalReactionConstant();

  int nExA = ExAList.size();
  int nEx  = ExList.size();

  // optional visualization control: pass "vis" as 3rd arg
  bool enableVis = (argc >= 3 && strcmp(argv[2], "vis") == 0);
  TApplication *app = nullptr;
  if(enableVis){
    app = new TApplication("anasenVis", &argc, argv);
  }

  // storage for tracks during simulation (for visualization)
  std::vector<TVector3> visTrackVertex, visTrackDir, visTrackHitPos;
  std::vector<std::pair<int,int>> visTrackWires;  // {anodeID, cathodeID}

  // create detector representation in memory
  ANASEN * anasen = new ANASEN();       // top-level detector object
  SX3 * sx3 = anasen->GetSX3();         // silicon array part
  PW * pw = anasen->GetPW();            // proportional wire chamber part

  // output file + tree
  TString saveFileName = "SimAnasen1.root";
  printf("\e[32m#################################### building Tree in %s\e[0m\n", saveFileName.Data());
  TFile * saveFile = new TFile(saveFileName, "recreate");
  TTree * tree = new TTree("tree", "tree");
   

  // beam and CM variables saved in tree
  double KEA;
  tree->Branch("beamKEA",     &KEA, "beamKEA/D");

  double thetaCM, phiCM;
  tree->Branch("thetaCM", &thetaCM, "thetaCM/D");
  tree->Branch("phiCM",     &phiCM, "phiCM/D");

  // outgoing particles in lab frame (light/heavy)
  double thetab, phib, Tb;
  double thetaB, phiB, TB;
  double dEb;
  double dEB;
  std::array<double, 2> T;
  tree->Branch("thetab", &thetab, "thetab/D"); // polar angle of light particle in lab frame
  tree->Branch("phib",     &phib, "phib/D"); // azimuthal angle of light particle in lab frame
  tree->Branch("Tb",         &Tb, "Tb/D"); // kinetic energy of light particle at vertex (before energy loss)
  tree->Branch("thetaB", &thetaB, "thetaB/D");
  tree->Branch("phiB",     &phiB, "phiB/D");
  tree->Branch("TB",         &TB, "TB/D"); // kinetic energy of heavy particle at vertex (before energy loss)
  tree->Branch("dEb",        &dEb, "dEb/D");
  tree->Branch("dEB",        &dEB, "dEB/D"); // placeholder for heavy particle energy loss, currently set equal to light particle loss for simplicity
  tree->Branch("T",          &T, "T/D");   // placeholder for true Q-value, currently set to 0 for simplicity

  // excitation state identifiers
  int ExAID;
  double ExA;
  tree->Branch("ExAID", &ExAID, "ExAID/I"); // projectile excitation state ID
  tree->Branch("ExA",     &ExA, "ExA/D"); // projectile excitation energy in MeV

  int ExID;
  double Ex;
  tree->Branch("ExID", &ExID, "ExID/I"); // target excitation state ID
  tree->Branch("Ex",     &Ex, "Ex/D"); // target excitation energy in MeV

  // true vertex position in target volume
  double vertexX, vertexY, vertexZ;
  tree->Branch("vX",     &vertexX, "VertexX/D"); // true vertex X position in mm
  tree->Branch("vY",     &vertexY, "VertexY/D"); // true vertex Y position in mm
  tree->Branch("vZ",     &vertexZ, "VertexZ/D"); // true vertex Z position in mm

  // reconstructed SX3 hit position
  double sx3X, sx3Y, sx3Z;
  tree->Branch("sx3X",     &sx3X, "sx3X/D"); // reconstructed X position from SX3 (with optional smearing)
  tree->Branch("sx3Y",     &sx3Y, "sx3Y/D"); // reconstructed Y position from SX3 (with optional smearing)
  tree->Branch("sx3Z",     &sx3Z, "sx3Z/D"); // reconstructed Z position from SX3 (with optional smearing)

  // PW nearest and next nearest wires
  int anodeID[2], cathodeID[2];
  tree->Branch("aID", anodeID, "anodeID/I"); // anodeID[0] is nearest anode wire, anodeID[1] is next nearest anode wire
  tree->Branch("cID", cathodeID, "cathodeID/I"); // cathodeID[0] is nearest cathode wire, cathodeID[1] is next nearest cathode wire

  // distances to nearest wires
  double anodeDist[2], cathodeDist[2];
  tree->Branch("aDist",   anodeDist, "anodeDist/D");
  tree->Branch("cDist",  cathodeDist, "cathodeDist/D");

  // SX3 channel assignment and Z fraction (depth) information
  int sx3ID, sx3Up, sx3Dn, sx3Bk;
  double sx3ZFrac;
  tree->Branch("sx3ID",   &sx3ID, "sx3ID/I");
  tree->Branch("sx3Up",   &sx3Up, "sx3Up/I");
  tree->Branch("sx3Dn",   &sx3Dn, "sx3Dn/I");
  tree->Branch("sx3Bk",   &sx3Bk, "sx3Bk/I");
  tree->Branch("sx3ZFrac", &sx3ZFrac, "sx3ZFrac/D");

  // reconstructed angles from PW track fit, method 1 and 2
  double reTheta, rePhi;
  tree->Branch("reTheta", &reTheta, "reconstucted_theta/D");
  tree->Branch("rePhi",     &rePhi, "reconstucted_phi/D");

  double reTheta1, rePhi1;
  tree->Branch("reTheta1", &reTheta1, "reconstucted_theta1/D");
  tree->Branch("rePhi1",     &rePhi1, "reconstucted_phi1/D");

  // reconstructed vertex Z from PW fit
  double z0;
  tree->Branch("z0",     &z0, "reconstucted_Z/D");

  TTree* tree2 = tree->CloneTree(0);
  tree2->SetName("tree2");

  //========timer
  TBenchmark clock;
  bool shown ;
  clock.Reset();
  clock.Start("timer");
  shown = false;
  int ELossTotal = 0;

  //================================= Calculate event loop
  for( int i = 0; i < numEvent ; i++){

    // randomly sample target/projectile excitations
    ExAID = gRandom->Integer(nExA);
    ExA = ExAList[ExAID];
    transfer.SetExA(ExA);

    ExID = gRandom->Integer(nEx);
    Ex = ExList[ExID];
    transfer.SetExB(Ex);

    // recalc kinematic constants for chosen states
    transfer.CalReactionConstant();

    // isotropic CM direction
    thetaCM = TMath::ACos(2 * gRandom->Rndm() - 1) ;
    phiCM   = (gRandom->Rndm() - 0.5) * TMath::TwoPi();

    //==== Calculate reaction kinematics in lab frame
    TLorentzVector * output = transfer.Event(thetaCM, phiCM); // returns array of outputs
    TLorentzVector Pb = output[2]; // light particle or product A
    TLorentzVector PB = output[3]; // heavy particle or product B

    thetab = Pb.Theta() * TMath::RadToDeg();
    thetaB = PB.Theta() * TMath::RadToDeg();

    Tb = (Pb.E() - Pb.M()); // kinetic energy of light particle at vertex (before energy loss) units of MeV
    TB = (PB.E() - PB.M());
    T[0] = Tb;
    T[1] = TB;
    //if (Tb < 1.5) {
    //  //skip event if light particle energy after loss is below detection threshold of 1.5 MeV
    //  continue;
    //}

    phib = Pb.Phi() * TMath::RadToDeg();
    phiB = PB.Phi() * TMath::RadToDeg();

    // vertex position in target volume
    vertexX = (vertexXRange[1]- vertexXRange[0])*gRandom->Rndm() +  vertexXRange[0];
    vertexY = (vertexYRange[1]- vertexYRange[0])*gRandom->Rndm() +  vertexYRange[0];
    vertexZ = (vertexZRange[1]- vertexZRange[0])*gRandom->Rndm() +  vertexZRange[0];

    TVector3 vertex(vertexX, vertexY, vertexZ);

    // set direction vector from lab angle
    TVector3 dir(1, 0, 0);
    dir.SetTheta(thetab * TMath::DegToRad());
    dir.SetPhi(phib * TMath::DegToRad());

    // run detector response models for PW and SX3
    pw->FindWireID(vertex, dir, false);
    sx3->FindSX3Pos(vertex, dir, false);

    PWHitInfo hitInfo = pw->GetHitInfo();

    anodeID[0]   = hitInfo.nearestWire.first; // nearest anode wire ID
    cathodeID[0] = hitInfo.nearestWire.second; // nearest cathode wire ID
    anodeID[1]   = hitInfo.nextNearestWire.first; // next nearest anode wire ID
    cathodeID[1] = hitInfo.nextNearestWire.second; // next nearest cathode wire ID

    anodeDist[1]   = hitInfo.nextNearestDist.first; // distance to next nearest anode wire
    cathodeDist[1] = hitInfo.nextNearestDist.second; // distance to next nearest cathode wire

    if(IsDeadAnode(anodeID[0])) continue;
    if(IsDeadCathode(cathodeID[0])) continue;

    // SX3 hit channel info and depth fraction
    sx3ID = sx3->GetID();

    if(IsDeadSX3(sx3ID)) continue;

    anodeDist[0]   = hitInfo.nearestDist.first; // distance to nearest anode wire
    cathodeDist[0] = hitInfo.nearestDist.second; // distance to nearest cathode wire

    if( sx3ID >= 0 ){
      sx3Up    = sx3->GetChUp();
      sx3Dn    = sx3->GetChDn();
      sx3Bk    = sx3->GetChBk();
      sx3ZFrac = sx3->GetZFrac();
      
      // apply intrinsic detector resolution to true SX3 hit position
      // for no smearing comment out and use GetHitPos();
      TVector3 hitPos = sx3->GetHitPosWithSigma(sigmaSX3_W, sigmaSX3_L);

      sx3X = hitPos.X();
      sx3Y = hitPos.Y();
      sx3Z = hitPos.Z();

      // store track data for visualization if enabled
      if(enableVis){
        visTrackVertex.push_back(vertex);
        visTrackDir.push_back(dir);
        visTrackHitPos.push_back(hitPos);
        visTrackWires.push_back({anodeID[0], cathodeID[0]});
      }
      // reconstruct track from PW readings + SX3 hit
      pw->CalTrack(hitPos, anodeID[0], cathodeID[0], false);
      reTheta = pw->GetTrackTheta() * TMath::RadToDeg();
      rePhi   = pw->GetTrackPhi() * TMath::RadToDeg();

      // alternative track algorithm with uncertainty parameters
      pw->CalTrack2(hitPos, hitInfo, sigmaPW_A, sigmaPW_C, false);
      reTheta1 = pw->GetTrackTheta() * TMath::RadToDeg();
      rePhi1   = pw->GetTrackPhi() * TMath::RadToDeg();

      z0 = pw->GetZ0();
      dEb = 0;
      dEB = 0;
      tree->Fill();

      //Energy loss
      double dl = (hitPos - vertex).Mag(); // path length in units of cm
      if (numEvent <= 100){
        //printf("Event %d: Ekin before loss = %f MeV, distance = %f cm\n", i, Tb, dl);
        //printf("Total T before loss: %f MeV\n", T);
      }
      double tb_temp = Tb;

      dEb = tb_temp - Tb; // total energy loss
      if (ReactionProductb.count({4, 2})){ // if light particle is alpha, use alpha energy loss table
        double x0b = invgAlpha->Eval(Tb);
        x0b = x0b + dl;
        Tb = elossAlpha->Eval(x0b);
      } else if (ReactionProductb.count({1, 1})){ // if light particle is proton, use proton energy loss table
        double x0b = invgProton->Eval(Tb);
        x0b = x0b + dl;
        Tb = elossProton->Eval(x0b);
      } else {
        // for other particle types, can add additional energy loss tables or use a generic approximation
        // for now, we will just apply a simple linear energy loss as a placeholder
        double dE_dx = 5; // MeV/cm, placeholder value for energy loss per unit length
        Tb = Tb - dE_dx * dl;
      }

      //if (Tb < 0) {
      //  Tb = TMath::QuietNaN();
      //}

      dEb = tb_temp - Tb; // total energy loss
      
      // fill tree2 with energy loss adjusted data
      //Fill T so it can make a histogram of both Tb and TB in root script
      T[0] = Tb;
      T[1] = 0; 
      //to plot both as one histogram in root, can use tree2->Draw("T(0)"); for light particle and tree2->Draw("T(1)") for heavy particle

      tree2->Fill();
      
      if (numEvent <= 10){
        //printf("Event %d: Tb after energy loss = %f MeV, energy loss = %f MeV\n", i, Tb, tb_temp - Tb);
      } //to give in scientific notation, use %e instead of %f in the printf format string. For example: printf("Event %d: Tb after energy loss = %e MeV, energy loss = %e MeV\n", i, Tb, tb_temp - Tb);
      ELossTotal += (tb_temp - Tb);

    }else{
      // no valid SX3 hit: mark clearly invalid
      sx3Up = -1;
      sx3Dn = -1;
      sx3Bk = -1;
      sx3ZFrac = TMath::QuietNaN();

      sx3X = TMath::QuietNaN();
      sx3Y = TMath::QuietNaN();
      sx3Z = TMath::QuietNaN();

      reTheta  = TMath::QuietNaN();
      rePhi    = TMath::QuietNaN();
      reTheta1 = TMath::QuietNaN();
      rePhi1   = TMath::QuietNaN();
      z0       = TMath::QuietNaN();
      dEb      = TMath::QuietNaN();
      dEB      = TMath::QuietNaN();
      //Tb = -12354567; // mark kinetic energy as invalid for no hit case
      // fill tree with original data (no energy loss for these events)
      //comment out tree fill for no hit case  
      //tree->Fill();
      //tree2->Fill();
    }

    //#################################################################### Timer
    // measure elapsed real time and print progress roughly every 10 sec
    clock.Stop("timer");
    Double_t time = clock.GetRealTime("timer");
    clock.Start("timer");

    if ( !shown ) {
      if (fmod(time, 10) < 1 ){
        printf( "%10d[%2d%%]| %8.2f sec | expect: %5.1f min \n", i, TMath::Nint((i+1)*100./numEvent), time , numEvent*time/(i+1)/60);
        shown = 1;
      }
    } else {
      if (fmod(time, 10) > 9 ){
        shown = 0;
      }
    }

  }

  // write results to ROOT file and close
  //tree->Write();
  //tree2->Write();
  tree->Write("", TObject::kOverwrite);
  tree2->Write("", TObject::kOverwrite);
  int count = tree->GetEntries();
  int count2 = tree2->GetEntries();
  saveFile->Close();

  printf("=============== done. saved as %s. tree entries: %d, tree2 entries: %d\n", saveFileName.Data(), count, count2);
  printf("Total energy loss across all events: %f MeV\n", (double)ELossTotal);
  printf("Average energy loss across events: %f MeV\n", (double)ELossTotal / count);

  if(enableVis){ // to enable visualization, run with 3rd argument "vis", e.g. "./anasenMC 1000 vis"
    printf("Displaying geometry with %zu tracks from simulation\n", visTrackVertex.size());
    
    // Build full geometry with all wires
    anasen->DrawAnasen(0, 23, 0, 23, -1, true);
    
    // Add all stored tracks to the geometry
    TGeoManager *geom = anasen->GetGeoManager();
    TGeoVolume *worldBox = anasen->GetWorldBox();
    
    if(geom && worldBox && visTrackVertex.size() > 0){
      int trackNodeID = 500;  // start node IDs for tracks
      
      for(size_t iTrack = 0; iTrack < visTrackVertex.size(); ++iTrack){
        TVector3 vertex = visTrackVertex[iTrack];
        TVector3 dir = visTrackDir[iTrack];
        TVector3 hitPos = visTrackHitPos[iTrack];
        
        double theta = dir.Theta() * TMath::RadToDeg();
        double phi = dir.Phi() * TMath::RadToDeg();
        
        // Add a line marker at the vertex
        TGeoVolume *startMarker = geom->MakeSphere("startMarker", 0, 0, 2.0);
        startMarker->SetLineColor(kBlack);
        worldBox->AddNode(startMarker, trackNodeID, 
                         new TGeoCombiTrans(vertex.X(), vertex.Y(), vertex.Z(), 
                                          new TGeoRotation("rot", 0, 0, 0)));
        trackNodeID++;
        
        // Add track line from vertex toward hit position
        TGeoVolume *trackLine = geom->MakeTube("trackLine", 0, 0, 0.08, 150.0);
        trackLine->SetLineColor(kBlue);
        worldBox->AddNode(trackLine, trackNodeID,
                         new TGeoCombiTrans(vertex.X(), vertex.Y(), vertex.Z(),
                                          new TGeoRotation("rotTrack", phi + 90, theta, 0)));
        trackNodeID++;
        
        // Add hit position marker
        TGeoVolume *hitMarker = geom->MakeSphere("hitMarker", 0, 0, 2.0);
        hitMarker->SetLineColor(kRed);
        worldBox->AddNode(hitMarker, trackNodeID,
                         new TGeoCombiTrans(hitPos.X(), hitPos.Y(), hitPos.Z(),
                                          new TGeoRotation("rotHit", 0, 0, 0)));
        trackNodeID++;
      }
      
      // Redraw geometry with all tracks
      geom->CloseGeometry();
      geom->SetVisLevel(4);
      worldBox->Draw("ogle");
    }
    
    if(app){
      printf("Entering ROOT event loop\n");
      app->Run();
    }
  }

  delete anasen;
  delete elossAlpha;
  delete elossProton;


  return 0;

}