import json
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

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


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


E1, x1, y1, z1 = four(df.mu1_pt.to_numpy(), df.mu1_eta.to_numpy(), df.mu1_phi.to_numpy())
E2, x2, y2, z2 = four(df.mu2_pt.to_numpy(), df.mu2_eta.to_numpy(), df.mu2_phi.to_numpy())
m2 = (E1 + E2) ** 2 - (x1 + x2) ** 2 - (y1 + y2) ** 2 - (z1 + z2) ** 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)
res = {
    "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.mu1_charge * df.mu2_charge < 0)).sum()),
}
json.dump(res, open("result.json", "w"), indent=2)
os.makedirs("figures", exist_ok=True)
fig, ax = plt.subplots(figsize=(7, 4.5))
ax.hist(m, bins=160, range=(40, 200), histtype="stepfilled", alpha=0.7)
ax.set_xlabel("Dimuon invariant mass m_ll [GeV]")
ax.set_ylabel("Events / 1 GeV")
ax.set_yscale("log")
fig.savefig("figures/m_ll.svg")
print(res)