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

df = pd.read_csv("data/cms_dimuon_2011.csv")
px, py, pz, E = df.px1 + df.px2, df.py1 + df.py2, df.pz1 + df.pz2, df.E1 + df.E2
pt_ll = np.hypot(px, py)
y_ll = 0.5 * np.log((E + pz) / (E - pz))
dphi = np.abs(df.phi1 - df.phi2)
dphi = np.where(dphi > np.pi, 2 * np.pi - dphi, dphi)
deta = df.eta1 - df.eta2
dr = np.sqrt(deta**2 + dphi**2)
out = pd.DataFrame({"Run": df.Run, "Event": df.Event, "pt_ll": pt_ll, "y_ll": y_ll, "delta_phi": dphi, "delta_r": dr})
out.to_csv("kinematics.csv", index=False)
res = {"n_events": int(len(df)), "mean_pt_ll": float(pt_ll.mean()), "mean_delta_r": float(dr.mean()),
       "frac_delta_phi_gt_2p5": float((dphi > 2.5).mean()), "max_abs_y_ll": float(np.abs(y_ll).max())}
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(pt_ll, bins=100, range=(0, 200), histtype="stepfilled", alpha=0.7)
ax.set_xlabel("Dimuon transverse momentum pT(ll) [GeV]")
ax.set_ylabel("Events / 2 GeV")
ax.set_yscale("log")
fig.savefig("figures/pt_ll.svg")
print(res)