# -------------------------------------------------------------------------------------
# skeleton_cms_dimuon.py
#
# Example PyROOT to read/analyse cms data files
# for analyzing X -> mu mu
# -------------------------------------------------------------------------------------

import ROOT
import math

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

MonteCarlo = True

# --- Chain setup ---
chain = ROOT.TChain("Events")
print("Running CMS Recorded Data...")
fileout = ROOT.TFile("cmsDimuon.root", "recreate")
chain.Add("/hepdata/CMS/DoubleMuon/*.root")

# --- Histogram ---
hm_dimuon = ROOT.TH1F("hm_dimuon",";dimuon mass m_{#mu#mu} [GeV]; Entries",30000, 0.25, 300)

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

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

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

    # Read branches as Python lists
    if chain.nMuon != 2: continue

    ch  = list(chain.Muon_charge)
    if ch[0] == ch[1]: continue

    pt  = list(chain.Muon_pt)
    eta = list(chain.Muon_eta)
    phi = list(chain.Muon_phi)
    mass= list(chain.Muon_mass)

    # Build four-vectors
    p1 = ROOT.TLorentzVector()
    p2 = ROOT.TLorentzVector()
    p1.SetPtEtaPhiM(pt[0], eta[0], phi[0], mass[0])
    p2.SetPtEtaPhiM(pt[1], eta[1], phi[1], mass[1])

    dimuon = p1 + p2
    m = dimuon.M()
    hm_dimuon.Fill(m)

# --- Save and plot ---
fileout.Write()
fileout.Close()
print("\nDone.")