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

OBS = ["met", "mt", "n_jets", "lead_jet_pt", "ht", "delta_phi"]
df = pd.read_csv("data/synth_sigbkg.csv")
sig, bkg = df[df.label == 1], df[df.label == 0]
n_s, n_b = len(sig), len(bkg)


def auc(x_s, x_b):
    r = rankdata(np.concatenate([x_s, x_b]))  # average ranks handle ties as 1/2
    return float((r[: len(x_s)].sum() - len(x_s) * (len(x_s) + 1) / 2) / (len(x_s) * len(x_b)))


aucs = {c: auc(sig[c].to_numpy(float), bkg[c].to_numpy(float)) for c in OBS}
ranked = sorted(OBS, key=lambda c: -aucs[c])
best = ranked[0]
res = {"auc": aucs, "features_ranked": ranked, "best_feature": best, "n_signal": int(n_s), "n_background": int(n_b),
       "excluded_columns": [c for c in df.columns if c not in OBS]}
json.dump(res, open("result.json", "w"), indent=2)
os.makedirs("figures", exist_ok=True)
lo, hi = float(df[best].min()), float(np.percentile(df[best], 99.5))
fig, ax = plt.subplots(figsize=(7, 4.5))
ax.hist(sig[best], bins=60, range=(lo, hi), density=True, histtype="step", label="signal")
ax.hist(bkg[best], bins=60, range=(lo, hi), density=True, histtype="step", label="background")
ax.set_xlabel(f"{best} [GeV]" if best not in ("n_jets", "delta_phi") else best)
ax.set_ylabel("Normalised events")
ax.legend()
fig.savefig("figures/best_feature.svg")
print(res)