# -------------------------------------------------------------------------------------
# skeleton_aleph_pi02gg.py
#
# Example PyROOT to read/analyse Aleph files
# for analysing pi0 -> g g events (for both Monte Carlo and Recorded data)
# -------------------------------------------------------------------------------------

import ROOT
import math

MonteCarlo = False

# --- Enable multi-threading ---
ROOT.EnableImplicitMT()

# --- Chain setup ---
Event = ROOT.TChain("EventInfo")
Gamma = ROOT.TChain("Gamma")
Track = ROOT.TChain("Track")
Truth = ROOT.TChain("Truth")

if MonteCarlo:
  print("Running Aleph Monte Carlo")
  fileout = ROOT.TFile("histoMC.root", "recreate")
  Event.Add("/hepdata/Aleph/AlephMC91.root")
  Gamma.Add("/hepdata/Aleph/AlephMC91.root")
  Track.Add("/hepdata/Aleph/AlephMC91.root")
  Truth.Add("/hepdata/Aleph/AlephMC91.root")
else:
  print("Running Aleph Real Data")
  fileout = ROOT.TFile("histoDA.root", "recreate")
  Event.Add("/hepdata/Aleph/AlephDA91.root")
  Gamma.Add("/hepdata/Aleph/AlephDA91.root")
  Track.Add("/hepdata/Aleph/AlephDA91.root")

# --- Histograms ---
hE_all   = ROOT.TH1F("hE_all", ";E [GeV];Entries", 100, 0, 5)
hmgg_all = ROOT.TH1F("hmgg_all", ";m({#gamma,#gamma}) [GeV];Entries", 100, 0, 1)
hmgg_sig = ROOT.TH1F("hmgg_sig", ";m({#gamma,#gamma}) [GeV];Entries", 100, 0, 1)
hmgg_bgr = ROOT.TH1F("hmgg_bgr", ";m({#gamma,#gamma}) [GeV];Entries", 100, 0, 1)

# --- Event loop ---
fraction = 0.02
nentries = int(fraction * Gamma.GetEntries())
print(f"Nevents: {nentries}")

for ev in range(nentries):
    Event.GetEntry(ev)
    Gamma.GetEntry(ev)
    if MonteCarlo:
       Truth.GetEntry(ev)

    # Select good hadronic events (e+ e- -> Z -> q qbar)
    isGoodEvent = list(Event.evt_goodevent)
    if isGoodEvent[0]==False: continue

    if ev % 5000 == 0:
        pct = math.ceil(100.0 * ev / nentries)
        print(f"{ev} / {nentries} [{pct}%]", end="\r", flush=True)

    # Read branches as Python lists
    px = list(Gamma.gam_px)
    py = list(Gamma.gam_py)
    pz = list(Gamma.gam_pz)
    tind = list(Gamma.gam_truthindex)

    if MonteCarlo:
       m1id  = list(Truth.tru_m1id)
       m1brc = list(Truth.tru_m1brc)

    # photon list
    pht = []

    # Build four-vectors
    for i in range(len(px)):
        g = ROOT.TLorentzVector()
        g.SetXYZM(px[i], py[i], pz[i], 0.0)
        if g.P()<1.0: continue
        hE_all.Fill(g.E())

        motherid, motherbrc = 0,0

        if MonteCarlo and tind[i] < 999:
              motherid  = m1id[tind[i]]
              motherbrc = m1brc[tind[i]]
        pht.append((g,motherid,motherbrc))

    if len(pht) > 1:
       for i in range(len(pht)-1):
           for j in range(i+1,len(pht)):
              pi0 = pht[i][0] + pht[j][0]; # add 4-vectors
              mass = pi0.M()
              hmgg_all.Fill(mass)
              if pht[i][1] == 7 and pht[i][2]==pht[j][2]: # pi0 signal
                  hmgg_sig.Fill(mass)
              else:
                  hmgg_bgr.Fill(mass)

print("\nEvent loop completed.")

hE_all.Write()
hmgg_all.Write()
hmgg_sig.Write()
hmgg_bgr.Write()
fileout.Close()

print("\nAll OK.")