"""
Interactive threshold explorer for two combined discriminators:
  1) mean absolute distance to median wire (spread)
  2) minimum number of active hits (nhits = sum of the wire profile)

Drag the two sliders to move each threshold and watch, live:
  - the signal (TP+FN) vs background (FP+TN) histograms for BOTH variables,
    with the accepted/rejected region shaded on each
  - the TP / FP / TN / FN counts from the COMBINED cut (spread >= thr_spread
    AND nhits >= thr_hits)
  - Efficiency, Purity, Accuracy, F1 for that combined cut
  - a marker showing where the combined cut lands relative to the
    single-variable "spread-only" ROC curve (reference curve)

Run it as a normal script (not headless) so matplotlib can open an
interactive window, e.g.:

    python interactive_median_distance_explorer.py

If you're on a remote/cluster machine without a display, either:
  - run with X11 forwarding (ssh -X), or
  - switch the backend to a notebook-friendly one (see NOTES at bottom).
"""

import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from sklearn.metrics import roc_curve, roc_auc_score

# ============================================================
# Load data
# ============================================================
PKL_PATH = "/nfs/fanae/user/jprado/LLP_Showers/shower-studies/wire_profile_run1.pkl"

with open(PKL_PATH, "rb") as f:
    data = pickle.load(f)

wireProfile = data["wireProfile"]


# ============================================================
# Recompute, per event: nhits and mean absolute distance to median wire
# (same definitions as in your script)
# ============================================================
def event_features(profiles):
    nhits_list = []
    mad_list = []
    for evt in profiles:
        active_wires = np.flatnonzero(evt)
        if len(active_wires) == 0:
            continue
        nhits = np.sum(evt)
        median_wire = np.median(active_wires)
        mad = np.mean(np.abs(active_wires - median_wire))
        nhits_list.append(nhits)
        mad_list.append(mad)
    return np.asarray(nhits_list, dtype=float), np.asarray(mad_list, dtype=float)


categories = ("TP", "FN", "FP", "TN")
feats_by_cat = {cat: event_features(wireProfile.get(cat, [])) for cat in categories}
# feats_by_cat[cat] = (nhits_array, mad_array)

signal_nhits = np.concatenate([feats_by_cat["TP"][0], feats_by_cat["FN"][0]])
signal_mad = np.concatenate([feats_by_cat["TP"][1], feats_by_cat["FN"][1]])

background_nhits = np.concatenate([feats_by_cat["FP"][0], feats_by_cat["TN"][0]])
background_mad = np.concatenate([feats_by_cat["FP"][1], feats_by_cat["TN"][1]])

# Large values are signal-like for both variables ->
# classify as "signal" when mad >= thr_spread AND nhits >= thr_hits
mad_min = min(signal_mad.min(), background_mad.min())
mad_max = max(signal_mad.max(), background_mad.max())

hits_min = min(signal_nhits.min(), background_nhits.min())
hits_max = max(signal_nhits.max(), background_nhits.max())

# Reference ROC curve using the spread variable alone (as before)
y_true = np.concatenate([np.ones(len(signal_mad)), np.zeros(len(background_mad))])
scores_mad = np.concatenate([signal_mad, background_mad])
fpr_ref, tpr_ref, _ = roc_curve(y_true, scores_mad)
auc_ref = roc_auc_score(y_true, scores_mad)


def metrics_at_thresholds(thr_spread, thr_hits):
    """Classify as 'signal' when mad >= thr_spread AND nhits >= thr_hits."""
    sig_mask = (signal_mad >= thr_spread) & (signal_nhits >= thr_hits)
    bkg_mask = (background_mad >= thr_spread) & (background_nhits >= thr_hits)

    TP = np.sum(sig_mask)
    FN = np.sum(~sig_mask)
    FP = np.sum(bkg_mask)
    TN = np.sum(~bkg_mask)

    efficiency = TP / (TP + FN) if (TP + FN) > 0 else 0.0
    purity = TP / (TP + FP) if (TP + FP) > 0 else 0.0
    accuracy = (TP + TN) / (TP + TN + FP + FN)
    F1 = 2 * purity * efficiency / (purity + efficiency) if (purity + efficiency) > 0 else 0.0
    fpr_pt = FP / (FP + TN) if (FP + TN) > 0 else 0.0
    tpr_pt = efficiency
    return TP, FP, TN, FN, efficiency, purity, accuracy, F1, fpr_pt, tpr_pt


# ============================================================
# Figure layout
# ============================================================
fig = plt.figure(figsize=(13, 8.5))
gs = fig.add_gridspec(2, 3, height_ratios=[3, 1])
ax_hist_mad = fig.add_subplot(gs[0, 0])
ax_hist_hits = fig.add_subplot(gs[0, 1])
ax_roc = fig.add_subplot(gs[0, 2])
ax_text = fig.add_subplot(gs[1, :])
ax_text.axis("off")

plt.subplots_adjust(bottom=0.28, wspace=0.35)

nbins = 100

# --- spread histogram ---
edges_mad = np.linspace(mad_min, mad_max, nbins + 1)
ax_hist_mad.hist(signal_mad, bins=edges_mad, density=True, alpha=0.5,
                  color="tab:green", label=f"Signal, n={len(signal_mad)}")
ax_hist_mad.hist(background_mad, bins=edges_mad, density=True, alpha=0.5,
                  color="tab:red", label=f"Background, n={len(background_mad)}")

thr_spread_line = ax_hist_mad.axvline(np.median(signal_mad), color="k", ls="--", lw=2)
shade_mad = ax_hist_mad.axvspan(np.median(signal_mad), mad_max, color="gray", alpha=0.15)

ax_hist_mad.set_xlabel(r"Mean abs. distance to median wire  $\frac{1}{N}\sum_i|w_i-w_{med}|$")
ax_hist_mad.set_ylabel("Probability density")
ax_hist_mad.set_title("Spread distribution")
ax_hist_mad.grid(alpha=0.3)
ax_hist_mad.legend(loc="upper right", fontsize=8)

# --- nhits histogram ---
edges_hits = np.linspace(hits_min, hits_max, nbins + 1)
ax_hist_hits.hist(signal_nhits, bins=edges_hits, density=True, alpha=0.5,
                   color="tab:green", label=f"Signal, n={len(signal_nhits)}")
ax_hist_hits.hist(background_nhits, bins=edges_hits, density=True, alpha=0.5,
                   color="tab:red", label=f"Background, n={len(background_nhits)}")

thr_hits_line = ax_hist_hits.axvline(np.median(signal_nhits), color="k", ls="--", lw=2)
shade_hits = ax_hist_hits.axvspan(np.median(signal_nhits), hits_max, color="gray", alpha=0.15)

ax_hist_hits.set_xlabel("Number of active hits")
ax_hist_hits.set_ylabel("Probability density")
ax_hist_hits.set_title("Hit-count distribution")
ax_hist_hits.grid(alpha=0.3)
ax_hist_hits.legend(loc="upper right", fontsize=8)

# --- ROC (spread-only, as reference) ---
ax_roc.plot(fpr_ref, tpr_ref, label=f"Spread-only ROC (AUC={auc_ref:.3f})")
ax_roc.plot([0, 1], [0, 1], "k--", alpha=0.6)
combo_point, = ax_roc.plot([], [], marker="*", markersize=18, color="red", zorder=5,
                            label="Combined cut (spread + hits)")
ax_roc.set_xlabel("False Positive Rate")
ax_roc.set_ylabel("True Positive Rate")
ax_roc.set_title("Operating point vs spread-only ROC")
ax_roc.grid(alpha=0.3)
ax_roc.legend(loc="lower right", fontsize=8)
ax_roc.set_xlim(0, 1)
ax_roc.set_ylim(0, 1.02)
ax_roc.set_aspect("equal", adjustable="box")

metrics_text = ax_text.text(0.02, 0.5, "", fontsize=12, family="monospace",
                             va="center", transform=ax_text.transAxes)

# --- Sliders ---
ax_slider_spread = plt.axes([0.15, 0.12, 0.7, 0.03])
ax_slider_hits = plt.axes([0.15, 0.05, 0.7, 0.03])

init_thr_spread = float(np.median(signal_mad))
init_thr_hits = float(np.median(signal_nhits))

slider_spread = Slider(ax_slider_spread, "Min spread", mad_min, mad_max,
                        valinit=init_thr_spread)
slider_hits = Slider(ax_slider_hits, "Min hits", hits_min, hits_max,
                      valinit=init_thr_hits)


def update(val):
    thr_spread = slider_spread.val
    thr_hits = slider_hits.val

    TP, FP, TN, FN, eff, pur, acc, F1, fpr_pt, tpr_pt = metrics_at_thresholds(
        thr_spread, thr_hits
    )

    thr_spread_line.set_xdata([thr_spread, thr_spread])
    thr_hits_line.set_xdata([thr_hits, thr_hits])

    global shade_mad, shade_hits
    shade_mad.remove()
    shade_mad = ax_hist_mad.axvspan(thr_spread, mad_max, color="gray", alpha=0.15)

    shade_hits.remove()
    shade_hits = ax_hist_hits.axvspan(thr_hits, hits_max, color="gray", alpha=0.15)

    combo_point.set_data([fpr_pt], [tpr_pt])

    metrics_text.set_text(
        f"Min spread = {thr_spread:.3f}    Min hits = {thr_hits:.1f}\n"
        f"TP={TP:5d}   FP={FP:5d}   TN={TN:5d}   FN={FN:5d}\n"
        f"Efficiency={eff:.3f}   Purity={pur:.3f}   Accuracy={acc:.3f}   F1={F1:.3f}"
    )
    fig.canvas.draw_idle()


slider_spread.on_changed(update)
slider_hits.on_changed(update)
update(None)  # initialize text/markers

plt.show()

# ============================================================
# NOTES
# ============================================================
# - If running on a headless cluster node with no display, either:
#     ssh -X your_cluster   (X11 forwarding), or
#     run this in a Jupyter notebook and swap Slider for ipywidgets:
#         %matplotlib widget
#       (requires `pip install ipympl`), keeping the same logic.
# - Classification rule: an event is called "signal" when
#     mad >= Min spread   AND   nhits >= Min hits
#   i.e. keep events with HIGHER spread and MORE active hits.
# - The ROC curve shown is for the spread variable alone (kept as a
#   reference); the red star shows where the current COMBINED two-variable
#   cut lands in FPR/TPR space, which won't sit exactly on that curve since
#   it's driven by two variables at once rather than one.