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

LO, HI, W = 80.0, 100.0, 0.5
df = pd.read_csv("data/cms_dimuon_2011.csv")
sel = (df.type1 == "G") & (df.type2 == "G") & (df.eta1.abs() < 2.4) & (df.eta2.abs() < 2.4) & (df.pt1 > 20) & (df.pt2 > 20) & (df.Q1 * df.Q2 < 0)
d = df[sel]
m = d["M"].to_numpy()
in_range = (m >= LO) & (m <= HI)
counts, edges = np.histogram(m, bins=40, range=(LO, HI))
x = 0.5 * (edges[:-1] + edges[1:])


def gauss(m, mu, sigma):
    return np.exp(-0.5 * ((m - mu) / sigma) ** 2) / (sigma * np.sqrt(2 * np.pi))


def expo(m, lam):
    return np.exp(-m / lam) / (lam * (np.exp(-LO / lam) - np.exp(-HI / lam)))


def model(m, mu, sigma, ns, nb, lam):
    return W * (ns * gauss(m, mu, sigma) + nb * expo(m, lam))


ok = counts > 0
n_in = int(in_range.sum())
p0 = [91.0, 2.0, 0.9 * n_in, 0.1 * n_in, 30.0]
popt, pcov = curve_fit(model, x[ok], counts[ok], p0=p0, sigma=np.sqrt(counts[ok]), absolute_sigma=True, maxfev=20000)
err = np.sqrt(np.diag(pcov))
chi2 = float((((counts[ok] - model(x[ok], *popt)) ** 2) / counts[ok]).sum())
res = {
    "n_total": int(len(df)), "n_selected": int(len(d)), "efficiency": float(len(d) / len(df)), "n_in_range": n_in,
    "fit": {"mu": float(popt[0]), "mu_err": float(err[0]), "sigma": float(abs(popt[1])), "sigma_err": float(err[1]), "n_sig": float(popt[2]),
            "n_sig_err": float(err[2]), "n_bkg": float(popt[3]), "lam": float(popt[4]), "chi2": chi2, "ndf": int(ok.sum() - 5)},
}
json.dump(res, open("result.json", "w"), indent=2)
os.makedirs("figures", exist_ok=True)
fig, ax = plt.subplots(figsize=(7, 4.5))
ax.errorbar(x, counts, yerr=np.sqrt(counts), fmt="o", ms=3, label="data")
xx = np.linspace(LO, HI, 400)
ax.plot(xx, model(xx, *popt), label="fit (Gaussian + exponential)")
ax.set_xlabel("Dimuon invariant mass M [GeV]")
ax.set_ylabel("Events / 0.5 GeV")
ax.legend()
fig.savefig("figures/zfit.svg")
f = res["fit"]
open("REPORT.md", "w").write(f"""# Z peak in CMS 2011 dimuon open data

## Selection
Both muons global, |eta| < 2.4, pT > 20 GeV, opposite charge: {res['n_selected']} of {res['n_total']} events (efficiency {res['efficiency']:.4f}); {n_in} in [80, 100] GeV.

## Fit
Gaussian + exponential, 40 bins of 0.5 GeV, least squares with sqrt(n) errors: mu = {f['mu']:.3f} +- {f['mu_err']:.3f} GeV, sigma = {f['sigma']:.3f} GeV, N_s = {f['n_sig']:.0f} +- {f['n_sig_err']:.0f}, chi2/ndf = {f['chi2']:.1f}/{f['ndf']}.

## Assumptions
- The Gaussian is a resolution-dominated approximation of the Z line shape; no Breit-Wigner convolution and no radiative tail are modelled, so chi2/ndf is expected to exceed 1.
- Bins with zero counts are excluded from the chi2 (none occur here).
- No trigger, isolation or vertex requirements beyond the listed cuts; no efficiency corrections are applied.
- The exponential background is normalised on the fit range only.
""")
print(res)