"""Dimuon invariant mass pipeline (v1 ntuple). Author: previous student."""
import json
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

MU = 0.1056583755  # GeV

df = pd.read_csv("data/synth_zmumu.csv")

# v1 ntuple stores MeV -> convert to GeV
for c in ["pt1", "pt2"]:
    df[c] = df[c] / 1000.0


def four_vector(pt, eta, phi):
    px = pt * np.cos(phi)
    py = pt * np.sin(phi)
    pz = pt * np.sin(eta)
    e = np.sqrt(px**2 + py**2 + pz**2 + MU**2)
    return e, px, py, pz


e1, px1, py1, pz1 = four_vector(df["pt1"], df["eta1"], df["phi1"])
e2, px2, py2, pz2 = four_vector(df["pt2"], df["eta2"], df["phi2"])
m2 = (e1 + e2) ** 2 - (px1 + px2) ** 2 - (py1 + py2) ** 2 - (pz1 + pz2) ** 2
m = np.sqrt(np.clip(m2, 0, None))

pd.DataFrame({"event_id": df["event_id"], "m_ll": m}).to_csv("mass.csv", index=False)
win = (m > 80) & (m < 100)
result = {
    "n_events": int(len(df)),
    "n_in_window": int(win.sum()),
    "mean_mass_in_window": float(m[win].mean()),
    "median_mass_all": float(np.median(m)),
    "n_opposite_charge_in_window": int((win & (df["q1"] * df["q2"] < 0)).sum()),
    "bugs_fixed": [],
}
json.dump(result, open("result.json", "w"), indent=2)

try:
    fig, ax = plt.subplots()
    ax.hist(m, bins=160, range=(40, 200))
    ax.set_xlabel("m_ll [GeV]")
    ax.set_ylabel("Events")
    fig.savefig("figures/m_ll.svg")
except Exception:
    pass  # plotting is optional
print("done", result)

# Ntuple format changelog

## v2 (current, `data/synth_zmumu.csv`)
- Muon columns renamed from `pt1, eta1, phi1, q1, pt2, eta2, phi2, q2` to `mu1_pt, mu1_eta, mu1_phi, mu1_charge, mu2_pt, mu2_eta, mu2_phi, mu2_charge`.
- All momenta are stored in **GeV** (v1 stored MeV).
- Added `run`. `event_id` is unchanged.

## v1 (retired)
- Momenta in MeV; short column names.

{
 "n_events": 50000,
 "n_in_window": 34483,
 "mean_mass_in_window": 90.96866737789078,
 "median_mass_all": 90.38002673106806,
 "n_opposite_charge_in_window": 34205,
 "bugs_fixed": [
  "column names updated to the v2 schema (mu1_pt, ...)",
  "removed the MeV->GeV division: v2 already stores GeV",
  "pz = pt*sinh(eta), not pt*sin(eta)",
  "create figures/ before saving and do not swallow the exception"
 ]
}