# -------------------------------------------------------------------------------------
# skeleton_atlas_H2gg.py
#
# Example PyROOT to read/analyse Atlas data files
# for analyzing H -> g g events
# -------------------------------------------------------------------------------------

import ROOT
import math

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

# --- Output file and histogram ---
fout  = ROOT.TFile("mgg.root", "recreate")
h_mgg = ROOT.TH1F("h_mgg",";m_{#gamma#gamma} [GeV]; Entries",240, 100, 160)

# --- Chain setup ---
chain = ROOT.TChain("analysis")
chain.Add("/hepdata/Atlas/GamGam/ODEO_*.root")

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

for ev in range(nentries):
    chain.GetEntry(ev)

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

    # Read branches as Python lists
    pt       = list(chain.photon_pt)
    eta      = list(chain.photon_eta)
    phi      = list(chain.photon_phi)
    e        = list(chain.photon_e)
    tightID  = list(chain.photon_isTightID)
    ptcone20 = list(chain.photon_ptcone20)

    # Crack region veto
    if 1.37 < abs(eta[0]) < 1.52: continue
    if 1.37 < abs(eta[1]) < 1.52: continue

    # Tight ID
    if not tightID[0]: continue
    if not tightID[1]: continue

    # Isolation
    if ptcone20[0] / pt[0] > 0.055: continue
    if ptcone20[1] / pt[1] > 0.055: continue

    # Energy thresholds
    emax = e[0]
    emin = e[1]
    if e[1] > e[0]:
       emax = e[1]
       emin = e[0]
    if emax < 50: continue
    if emin < 30: continue

    # Build four-vectors
    g1 = ROOT.TLorentzVector()
    g2 = ROOT.TLorentzVector()
    g1.SetPtEtaPhiE(pt[0], eta[0], phi[0], e[0])
    g2.SetPtEtaPhiE(pt[1], eta[1], phi[1], e[1])

    gg = g1 + g2
    m  = gg.M()

    # Relative pT cuts
    if pt[0] / m < 0.35: continue
    if pt[1] / m < 0.35: continue

    h_mgg.Fill(m)

# --- Save and plot ---
fout.Write("h_mgg")

fout.Close()
print("\nAtlas merged data plotted.")