#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 <algorithm>
#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 "ClassQQQ.h"      // QQQ detector model class
#include "anasen_anode_cathode_hyperboloids.h"
#include "EnergyLoss.h"    // energy loss lookup between two positions in a medium
#include "HistPlotter.h"
#include "AutoHist2D.h"    // auto-ranged, auto-binned 2D histograms written alongside tree1
#include "Kinematics.h"
#include <stdio.h>
#include <stdlib.h>
#include <csignal>
#include <set>
#include <tuple>
#include "TLegend.h"
#include "TH1D.h"
#include "TObjArray.h"
#include "TBranch.h"
#include <iostream>
#include <fstream>
bool quit=false;
void handler(int sig){
  quit=true;
  printf("Caught signal %d, quitting gracefully...\n", sig);
}

//======== 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 std::string& filename) {
  TGraph* g = new TGraph(filename.c_str(), "%lg %lg");
  return g;
}

// Loads column 2 (Energy_MeV) vs column 4 (Sigma_x_cm) from an E_vs_x_*.dat table
TGraph* LoadSigmaXVsEnergy(const std::string& filename) {
  TGraph* g = new TGraph(filename.c_str(), "%*lg %lg %*lg %lg");
  g->Sort(); // TGraph::Eval requires ascending x (Energy_MeV)
  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 = {
      //0, 2, 4, 5, 6, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23
    }; // add dead SX3 IDs here, 0-23 1,7,9,3
    return dead.count(id);
}

bool IsDeadSX3FrontUpChannel(int sx3ID, int chUp){
  static std::set<std::pair<int, int>> dead = {
    // {sx3ID, front-up-channel}
  };
  return dead.count({sx3ID, chUp});
}

bool IsDeadSX3FrontDnChannel(int sx3ID, int chDn){
  static std::set<std::pair<int, int>> dead = {
    // {sx3ID, front-down-channel}
  };
  return dead.count({sx3ID, chDn});
}

bool IsDeadSX3BackChannel(int sx3ID, int chBk){
  static std::set<std::pair<int, int>> dead = {
    //{1, 10}
    // {sx3ID, back-channel} 
  };
  return dead.count({sx3ID, chBk});
}

bool IsDeadSX3ChannelCombo(int sx3ID, int chUp, int chDn, int chBk){
  return IsDeadSX3FrontUpChannel(sx3ID, chUp)
    || IsDeadSX3FrontDnChannel(sx3ID, chDn)
    || IsDeadSX3BackChannel(sx3ID, chBk);
}

// Simulate sequential two-body decay of an unstable parent in its rest frame.
TLorentzVector SimulateSequentialDecay(const TLorentzVector &parent,
                                       int daughterA, int daughterZ,
                                       int ejectA, int ejectZ,
                                       TLorentzVector &ejectileOut){
  Isotope daughter(daughterA, daughterZ);
  Isotope ejectile(ejectA, ejectZ);

  double M = parent.M();
  double mD = daughter.Mass;
  double mE = ejectile.Mass;

  double sqM = M * M;
  double sum = mD + mE;
  double diff = mD - mE;
  double p2 = (sqM - sum*sum) * (sqM - diff*diff) / (4.0 * sqM); // two-body decay momentum squared
  if( p2 < 0 ) p2 = 0; // handle unphysical case where parent mass is less than sum of daughter and ejectile masses
  double p = TMath::Sqrt(p2); // two-body decay momentum

  double cosTheta = 2.0 * gRandom->Rndm() - 1.0; // isotropic decay in parent rest frame
  double theta = TMath::ACos(cosTheta); // polar angle of daughter in parent rest frame
  double phi = gRandom->Rndm() * TMath::TwoPi(); // azimuthal angle of daughter in parent rest frame

  TVector3 v; // momentum vector of daughter in parent rest frame
  v.SetMagThetaPhi(p, theta, phi); // daughter momentum in parent rest frame

  TLorentzVector daughterLab; // daughter 4-vector in lab frame, initialized with momentum from decay and mass of daughter
  daughterLab.SetVectM(v, mD); // set daughter 4-vector in parent rest frame, then boost to lab frame

  TLorentzVector ejectileLab; // ejectile 4-vector in lab frame, initialized with momentum opposite to daughter and mass of ejectile
  ejectileLab.SetVectM(-v, mE); //  set ejectile 4-vector in parent rest frame, then boost to lab frame

  TVector3 boost = parent.BoostVector(); // boost vector to go from parent rest frame to lab frame
  daughterLab.Boost(boost); // boost daughter to lab frame
  ejectileLab.Boost(boost); // boost ejectile to lab frame

  ejectileOut = ejectileLab; // return ejectile in lab frame
  return daughterLab;
}

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]);
  TransferReaction transfer;

  // Register signal handler
  std::signal(SIGINT, handler);

  //To set beam energy loss, use energy loss app, and create table with target isotope, set Initial beam energy as max energy 
  transfer.SetA(27, 13, 0);            // 22Mg projectile
  transfer.Seta(4, 2);                 // 4He target
  transfer.Setb(1, 1);                 // outgoing proton from the primary transfer
  transfer.SetB(30, 14);               // 30Si* heavy product
  double beamE = 56.1;
  const ReactionConfig reactionConfig = transfer.GetRectionConfig();
  const double beamA = reactionConfig.beamA;                   // mass number of 14N beam

  // Excited state lists (projectile and heavy-product excitation states)
  std::vector<float> ExAList = {0};              // Beam excited energy
  std::vector<float> ExList = {0.0, 2.2, 3.4, 6.0};             // Heavy product excited energy

  const int kMBeam = reactionConfig.beamA; // mass number of beam
  const int kMTarget = reactionConfig.targetA; // mass number of target
  const int kMLight = reactionConfig.recoilLightA; // mass number of light ejectile
  const int kMHeavy = reactionConfig.recoilHeavyA; // mass number of heavy product

  //const int kZBeam = reactionConfig.beamZ; // atomic number of beam
  //const int kZTarget = reactionConfig.targetZ; // atomic number of target
  //const int kZLight = reactionConfig.recoilLightZ; // atomic number of light ejectile
  //const int kZHeavy = reactionConfig.recoilHeavyZ; // atomic number of heavy product
  //bool enableSequentialDecay = false; // turning to false to disable sequential decay for now, can be set to true to enable
  //const int decayDaughterA = 20;
  //const int decayDaughterZ = 10;
  //const int decayEjectA = 1;
  //const int decayEjectZ = 1;

  std::string b;
  if (reactionConfig.recoilLightA == 1) {
    b = "proton";
  } else if (reactionConfig.recoilLightA == 2) {
    b = "deuteron";
  } else if (reactionConfig.recoilLightA == 3) {
    b = "triton";
  } else if (reactionConfig.recoilLightA == 4) {
    b = "alpha";
  }

  //TGraph* elossLight = LoadELoss("../ELoss/HeLoss/E_vs_x_" + b + ".dat");
  // define vertex position uniform distribution ranges (mm)
  double vertexXRange[2] = {   -5, 5}; // mm - 5, 5
  double vertexYRange[2] = {   -5, 5}; // -5, 5
  double vertexZRange[2] = { -174.3, 174.3}; // -174.3, 174.3 (full length of gas volume, centered at 0)
  const double beamEntranceZ = -280 - 174.3; //vertexZRange[0]; // mm, assumed beam entrance into the gas
  
  TGraph* elossBeam = LoadELoss("../ELoss/HeLoss/E_vs_x_Al-27.dat"); // x = path length (cm), y = beam energy (MeV)
  TGraph* sigmaXBeam = LoadSigmaXVsEnergy("../ELoss/HeLoss/E_vs_x_Al-27.dat"); // x = beam energy (MeV), y = distance straggle sigma_x (cm)
  // Build a temporary inverse (energy -> path) to locate the path at beamE.
  TGraph* elossBeamInverseRaw = new TGraph(elossBeam->GetN());
  for( int p = 0; p < elossBeam->GetN(); p++ ){
    double x, y;
    elossBeam->GetPoint(p, x, y);
    elossBeamInverseRaw->SetPoint(p, y, x);
  }
  elossBeamInverseRaw->Sort(); // TGraph::Eval requires ascending x

  const double pathAtBeamE = elossBeamInverseRaw->Eval(beamE);

  // Keep only energies <= beamE and shift path so beamE corresponds to x = 0.
  TGraph* elossBeamFiltered = new TGraph();
  for( int p = 0; p < elossBeam->GetN(); p++ ){
    double x, y;
    elossBeam->GetPoint(p, x, y);
    if( y <= beamE ){
      const int n = elossBeamFiltered->GetN();
      elossBeamFiltered->SetPoint(n, x - pathAtBeamE, y);
    }
  }

  TGraph* elossBeamInverse = new TGraph(elossBeamFiltered->GetN());
  for( int p = 0; p < elossBeamFiltered->GetN(); p++ ){
    double x, y;
    elossBeamFiltered->GetPoint(p, x, y);
    elossBeamInverse->SetPoint(p, y, x);
  }
  elossBeamInverse->Sort(); // TGraph::Eval requires ascending x

  // detector resolution / uncertainty parameters
  double sigmaSX3_W = 0; // mm, if < 0 use mid-point (no spread in SX3 horizontal dimension)
  double sigmaSX3_L = 0;  // 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
  QQQ * qqq = anasen->GetQQQ();         // optional QQQ detector part, not used in this simulation but can be enabled for visualization

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

  // beam and CM variables saved in tree
  double eventID;
  double KEA;
  double beamPath_cm;
  int MBeamOut;
  int MTargetOut;
  int MLightOut;
  int MHeavyOut;
  int ZBeamOut;
  int ZTargetOut;
  int ZLightOut;
  int ZHeavyOut;
  double beamEnergy;
  tree1->Branch("eventID", &eventID, "eventID/D");
  tree1->Branch("beamKEA",      &KEA, "beamKEA/D");
  tree1->Branch("beamPath_cm",  &beamPath_cm, "beamPath_cm/D");
  tree1->Branch("beamEnergy",   &beamEnergy, "beamEnergy/D");
  tree1->Branch("MBeam", &MBeamOut, "MBeam/I");
  tree1->Branch("MTarget", &MTargetOut, "MTarget/I");
  tree1->Branch("MLight", &MLightOut, "MLight/I");
  tree1->Branch("MHeavy", &MHeavyOut, "MHeavy/I");
  tree1->Branch("ZBeam", &ZBeamOut, "ZBeam/I");
  tree1->Branch("ZTarget", &ZTargetOut, "ZTarget/I");
  tree1->Branch("ZLight", &ZLightOut, "ZLight/I");
  tree1->Branch("ZHeavy", &ZHeavyOut, "ZHeavy/I");

  // constant reaction mass numbers stored in every event entry
  MBeamOut = kMBeam;
  MTargetOut = kMTarget;
  MLightOut = kMLight;
  MHeavyOut = kMHeavy;

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

  // outgoing particles in lab frame (light/heavy)
  double thetab, phib, Tb, qqqTb, sx3Tb;
  double thetaB, phiB, TB, qqqTB, sx3TB;
  std::array<double, 2> T;
  tree1->Branch("thetab", &thetab, "thetab/D"); // polar angle of light particle in lab frame
  tree1->Branch("phib",     &phib, "phib/D"); // azimuthal angle of light particle in lab frame
  tree1->Branch("Tb",         &Tb, "Tb/D"); // kinetic energy of light particle at vertex (before energy loss)
  tree1->Branch("thetaB", &thetaB, "thetaB/D");
  tree1->Branch("phiB",     &phiB, "phiB/D");
  tree1->Branch("TB",         &TB, "TB/D"); // kinetic energy of heavy particle at vertex (before energy loss)
  tree1->Branch("T",          &T, "T/D");   // placeholder for true Q-value, currently set to 0 for simplicity
  tree1->Branch("qqqTb", &qqqTb, "qqqTb/D"); // kinetic energy of light particle at vertex (before energy loss) for events where the light particle hits the QQQ, currently set to 0 for simplicity
  tree1->Branch("qqqTB", &qqqTB, "qqqTB/D"); // kinetic energy of heavy particle at vertex (before energy loss) for events where the light particle hits the QQQ, currently set to 0 for simplicity
  tree1->Branch("sx3Tb", &sx3Tb, "sx3Tb/D"); // kinetic energy of light particle at vertex (before energy loss) for events where the light particle hits the SX3, currently set to 0 for simplicity
  tree1->Branch("sx3TB", &sx3TB, "sx3TB/D"); // kinetic energy of heavy particle at vertex (before energy loss) for events where the light particle hits the SX3, currently set to 0 for simplicity

  double Esx3, Eqqq, Edet;
  tree1->Branch("Esx3", &Esx3, "Esx3/D");
  tree1->Branch("Eqqq", &Eqqq, "Eqqq/D");
  tree1->Branch("Edet", &Edet, "Edet/D");

  double Eanode, Ecathode, EPC;
  tree1->Branch("Eanode", &Eanode, "Eanode/D");
  tree1->Branch("Ecathode", &Ecathode, "Ecathode/D");
  tree1->Branch("EPC", &EPC, "EPC/D");

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

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

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

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

  double qqqX, qqqY, qqqZ;
  tree1->Branch("qqqX",     &qqqX, "qqqX/D"); // reconstructed X position from QQQ (with optional smearing) in mm
  tree1->Branch("qqqY",     &qqqY, "qqqY/D"); // reconstructed Y position from QQQ (with optional smearing)
  tree1->Branch("qqqZ",     &qqqZ, "qqqZ/D"); // reconstructed Z position from QQQ (with optional smearing)

  double detX, detY, detZ;
  tree1->Branch("detX",     &detX, "detX/D");
  tree1->Branch("detY",     &detY, "detY/D");
  tree1->Branch("detZ",     &detZ, "detZ/D");

  double aX, aY, aZ;
  tree1->Branch("aX",     &aX, "aX/D");
  tree1->Branch("aY",     &aY, "aY/D");
  tree1->Branch("aZ",     &aZ, "aZ/D");

  double cX, cY, cZ;
  tree1->Branch("cX",     &cX, "cX/D");
  tree1->Branch("cY",     &cY, "cY/D");
  tree1->Branch("cZ",     &cZ, "cZ/D");

  double dl;
  tree1->Branch("dl", &dl, "dl/D");
  
  // PW nearest and next nearest wires
  int anodeID[2], cathodeID[2];
  tree1->Branch("aID", anodeID, "anodeID/I"); // anodeID[0] is nearest anode wire, anodeID[1] is next nearest anode wire
  tree1->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];
  tree1->Branch("aDist",   anodeDist, "anodeDist/D");
  tree1->Branch("cDist",  cathodeDist, "cathodeDist/D");

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

  tree1->Branch("qqqID",   &qqqID, "qqqID/I");
  tree1->Branch("qqqUp",   &qqqUp, "qqqUp/I");
  tree1->Branch("qqqBk",   &qqqBk, "qqqBk/I");

  double EBeam_Kin_gs=NAN, EBeam_Kin_2_2=NAN, EBeam_Kin_3_4=NAN, Ex_recon=NAN
    ;
  tree1->Branch("EBeam_Kin", &EBeam_Kin_gs, "EBeam_Kin/D");
  tree1->Branch("EBeam_Kin_2.2", &EBeam_Kin_2_2, "EBeam_Kin_2.2/D");
  tree1->Branch("EBeam_Kin_3.4", &EBeam_Kin_3_4, "EBeam_Kin_3.4/D");
  tree1->Branch("Ex_recon", &Ex_recon, "Ex_recon/D");

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

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

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

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

  //================================= Calculate event loop
  for( int i = 0; i < numEvent ; i++){
    if(quit) break; // exit gracefully if signal Ctrl+C received
    // randomly sample target/projectile excitations

    eventID = i;

    thetaCM = TMath::QuietNaN();
    phiCM = TMath::QuietNaN();
    thetab = TMath::QuietNaN();
    phib = TMath::QuietNaN();
    Tb = TMath::QuietNaN();
    thetaB = TMath::QuietNaN();
    phiB = TMath::QuietNaN();
    TB = TMath::QuietNaN();
    T = {TMath::QuietNaN(), TMath::QuietNaN()};
    qqqTb = TMath::QuietNaN();
    qqqTB = TMath::QuietNaN();
    sx3Tb = TMath::QuietNaN();
    sx3TB = TMath::QuietNaN();
    Esx3 = TMath::QuietNaN();
    Eqqq = TMath::QuietNaN();
    Eanode = TMath::QuietNaN();
    Ecathode = TMath::QuietNaN();
    EPC = TMath::QuietNaN();
    sx3X = TMath::QuietNaN();
    sx3Y = TMath::QuietNaN();
    sx3Z = TMath::QuietNaN();
    qqqX = TMath::QuietNaN();
    qqqY = TMath::QuietNaN();
    qqqZ = TMath::QuietNaN();
    detX = TMath::QuietNaN();
    detY = TMath::QuietNaN();
    detZ = TMath::QuietNaN();
    sx3ID = -1;
    sx3Up = -1;
    sx3Dn = -1;
    sx3Bk = -1;
    qqqID = -1;

    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();

    // vertex position in target volume
    beamEnergy = gRandom->Uniform(0, beamE); // MeV, sample beam energy at vertex from uniform distribution between 0 and initial beam energy
    KEA = beamEnergy / beamA;
    beamPath_cm = elossBeamInverse->Eval(beamEnergy); // beamE maps to x=0 after path shift
    vertexZ = beamEntranceZ + beamPath_cm * 10.0; // cm -> mm

    // transverse sampling range from the beam's distance straggle at this energy
    const double sigmaX_mm = std::max(0.0, sigmaXBeam->Eval(beamEnergy)) * 100.0; // cm -> mm
    //vertexX = 2.0 * sigmaX_mm * gRandom->Rndm() - sigmaX_mm;
    //vertexY = 2.0 * sigmaX_mm * gRandom->Rndm() - sigmaX_mm;
    double vertexRangeX = std::sqrt((vertexXRange[1] * vertexXRange[1]) + (sigmaX_mm * sigmaX_mm));
    double vertexRangeY = std::sqrt((vertexYRange[1] * vertexYRange[1]) + (sigmaX_mm * sigmaX_mm));
    //std::cout << "vertexRangeX: " << vertexRangeX << ", vertexRangeY: " << vertexRangeY << std::endl;
    vertexX = gRandom->Gaus(0, vertexRangeX); // mean and standard deviation
    vertexY = gRandom->Gaus(0, vertexRangeY); // mean and standard deviation

    //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);

    // compute beam energy at the event vertex from the gas path length
    
    //beamPath_cm = TVector3(vertexZ - beamEntranceZ, vertexX, vertexY).Mag() * 0.1;
    //beamDistance = vertexZ - beamEntranceZ;
    /*
    if( beamPath_cm < 0 ) beamPath_cm = 0;
    beamEnergy = elossBeam->Eval(beamPath_cm); // MeV
    double beamEnergyLoss = elossBeam->Eval(0.0) - beamEnergy;
    KEA = beamEnergy / beamA;*/
    //KEA = gRandom->Uniform(0, beamE);

    transfer.SetIncidentEnergyAngle(KEA, 0, 0);
    transfer.CalReactionConstant();

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

    //==== Calculate reaction kinematics in lab frame for the primary transfer
    TLorentzVector * output = transfer.Event(thetaCM, phiCM); // returns array of outputs
    TLorentzVector Pb = output[2]; // primary proton from transfer
    TLorentzVector PB = output[3]; // excited 21Na* heavy product

    thetab = Pb.Theta() * TMath::RadToDeg();
    Tb = (Pb.E() - Pb.M()); // kinetic energy of the light proton from the primary transfer
    thetaB = PB.Theta() * TMath::RadToDeg();
    TB = (PB.E() - PB.M());
    phib = Pb.Phi() * TMath::RadToDeg();
    phiB = PB.Phi() * TMath::RadToDeg();
    T[0] = Tb;
    T[1] = TB;

    delete [] output;

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

    qqq->Clear();
    sx3->Clear();
    // run detector response models for PW and SX3
    pw->FindWireID(vertex, dir, false);
    sx3->FindSX3Pos(vertex, dir, false);
    qqq->FindQQQPos(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])) anodeID[0] = -1; // mark as no hit if anode is dead
    if(IsDeadCathode(cathodeID[0])) cathodeID[0] = -1; // mark as no hit if cathode is dead

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

    if(IsDeadSX3(sx3ID)) sx3ID = -1; // mark as no hit if SX3 is dead

    anodeDist[0]   = hitInfo.nearestDist.first; // distance to nearest anode wire
    cathodeDist[0] = hitInfo.nearestDist.second; // distance to nearest cathode wire
    
    //start HERE
    if( sx3ID >= 0 ){

      sx3Up    = sx3->GetChUp();
      sx3Dn    = sx3->GetChDn();
      sx3Bk    = sx3->GetChBk();
      if(IsDeadSX3ChannelCombo(sx3ID, sx3Up, sx3Dn, sx3Bk)) sx3Up = -1, sx3Dn = -1, sx3Bk = -1; // mark as no hit if any SX3 channel is dead
      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();

      pw->CalTrack(hitPos, anodeID[0], cathodeID[0], false);

      const TVector3 trackDir = pw->GetTrackVec();
      const double transverseDirection2 =
          trackDir.X() * trackDir.X() + trackDir.Y() * trackDir.Y();

      if (transverseDirection2 > 0.0) {
        /*const double pathToRhoMin =
            -(hitPos.X() * trackDir.X() + hitPos.Y() * trackDir.Y())
            / transverseDirection2;*/

        //const TVector3 rhoMin = hitPos + pathToRhoMin * trackDir;
        const TVector3 rhoMin = vertex;

        // Original two-wire-layer model, kept for reference:
        auto [cathodeIntersection, anodeIntersection, pcPathLengthCm] =
            find_PC_PathLength(rhoMin, hitPos);
        
        // The helper returns cathode first, then anode.
        aX = anodeIntersection.X();
        aY = anodeIntersection.Y();
        aZ = anodeIntersection.Z();
        
        cX = cathodeIntersection.X();
        cY = cathodeIntersection.Y();
        cZ = cathodeIntersection.Z();
      }

      //Energy loss calculations
      double distance_sx3;
      Esx3 = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
               sx3X, sx3Y, sx3Z,
               b, "He", Tb,
           distance_sx3);
      dl = distance_sx3;

      double distance_A;
      Eanode = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
               aX, aY, aZ,
               b, "He", Tb,
           distance_A);

      double distance_C;
      Ecathode = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
               cX, cY, cZ,
               b, "He", Tb,
           distance_C);

      if (Esx3 <= 0 || Eanode <= 0 || Ecathode <= 0) {
        Esx3 = NAN;
        beamEnergy = NAN;
        Ex = NAN;
        //continue;
      }
      sx3Tb = Tb; // for simplicity, using the same kinetic energy for SX3 hit events, can be modified to simulate energy loss if desired
      sx3TB = TB;
      Edet = Esx3;
      EPC = Eanode - Ecathode;
      AutoHist2D::Fill("beamEnergy_vs_vZ", vertexZ / 10, beamEnergy, "vZ (cm)", "beamEnergy (MeV)");
      AutoHist2D::Fill("EPC x sin(theta) vs Esx3", Esx3, EPC * sin(thetab * TMath::DegToRad()), "Esx3 (MeV)", "EPC x sin(theta) (MeV)");
      //tree1->Fill();
      
      //Kinematics aakin_27Al(26.981538408,4.00260325413,4.0026035413,26.981538408,beam_energy_at_vertex/26.981538408); //m3 is alpha
      Kinematics apkin_27Al(26.981538408,4.00260325413,1.00782503224,29.973770136,beamEnergy/26.981538408); //m3 is proton
      //Kinematics apkin_27Al(1.00782503224,4.00260325413,4.00260325413,1.00782503224,beamEnergy/1.00782503224); //m3 is proton
      Ex_recon = apkin_27Al.getExc(Esx3, thetab);
      EBeam_Kin_gs = apkin_27Al.getEbeam_givenQ(Esx3, 0.0, thetab);
      EBeam_Kin_2_2 = apkin_27Al.getEbeam_givenQ(Esx3, 2.2, thetab);
      EBeam_Kin_3_4 = apkin_27Al.getEbeam_givenQ(Esx3, 3.4, thetab);
      //std::cout << EBeam_Kin << std::endl;

    }else if (qqqID >= 0){

      TVector3 hitPos = qqq->GetHitPos();

      qqqX = hitPos.X();
      qqqY = hitPos.Y();
      qqqZ = hitPos.Z();

        auto [cathodeIntersection, anodeIntersection, pcPathLengthCm] =
          find_PC_PathLength(vertex, hitPos);
        aX = anodeIntersection.X();
        aY = anodeIntersection.Y();
        aZ = anodeIntersection.Z();
        cX = cathodeIntersection.X();
        cY = cathodeIntersection.Y();
        cZ = cathodeIntersection.Z();

      if(enableVis){
        visTrackVertex.push_back(vertex);
        visTrackDir.push_back(dir);
        visTrackHitPos.push_back(hitPos);
        //visTrackWires.push_back({anodeID[0], cathodeID[0]});
      }

      //Energy loss calculations

      double distance_qqq;
      Eqqq = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
                           qqqX, qqqY, qqqZ,
                           b, "He",
               Tb, distance_qqq);
      dl = distance_qqq;
      double distance_A;
      Eanode = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
                           aX, aY, aZ,
                           b, "He",
               Tb, distance_A);
      double distance_C;
      Ecathode = CalculateEnergyLoss(vertexX, vertexY, vertexZ,
                           cX, cY, cZ,
                           b, "He",
               Tb, distance_C);

      if (Eqqq <= 0 || Eanode <= 0 || Ecathode <= 0) {
          Eqqq = NAN;
          Ex = NAN;
          //continue;
      }

      qqqTb = Tb; // for simplicity, using the same kinetic energy for QQQ hit events, can be modified to simulate energy loss if desired
      qqqTB = TB;
      Edet = Eqqq;
      EPC = Eanode - Ecathode;
      AutoHist2D::Fill("beamEnergy_vs_vZ", vertexZ / 10, beamEnergy, "vZ (cm)", "beamEnergy (MeV)");
      AutoHist2D::Fill("EPC x sin(theta) vs Eqqq", Eqqq, EPC * sin(thetab * TMath::DegToRad()), "Eqqq (MeV)", "EPC x sin(theta) (MeV)");

      beamEnergy = TMath::QuietNaN(); // mark beam energy as invalid for QQQ hit case
      //tree1->Fill();

    }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();
      Tb       = TMath::QuietNaN(); // mark kinetic energy as invalid for no hit case
      TB       = TMath::QuietNaN();
      sx3Tb     = TMath::QuietNaN();
      sx3TB     = TMath::QuietNaN();
      qqqTb     = TMath::QuietNaN();
      qqqTB     = TMath::QuietNaN();
      // fill tree with original data (no energy loss for these events)
      //comment out tree fill for no hit case  
      //tree1->Fill();
    }
    tree1->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
  AutoHist2D::WriteAll(); // books, fills and writes all registered 2D histograms into saveFile
  tree1->Write("", TObject::kOverwrite);
  int count1 = tree1->GetEntries();
  //int count3 = tree3->GetEntries();
  saveFile->Close();

  printf("=============== done. saved as %s. tree1 entries: %d\n", saveFileName.Data(), count1);

  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;


  return 0;

}