"""
Interactive threshold explorer for a two-parameter, "number of loud wires"
discriminator:

  1) Hit threshold (per wire): a wire counts as "loud" if its hit count is
     >= this value.
  2) Wire-count threshold: the event is classified as signal if the number
     of "loud" wires is >= this value.

For each event: n_loud_wires(evt, hit_thr) = sum(evt >= hit_thr)
Classification rule: signal if n_loud_wires >= wire_count_thr

Drag either slider and watch, live:
  - the distribution of "number of loud wires" for signal (TP+FN) vs
    background (FP+TN), recomputed for the current hit threshold
  - the TP / FP / TN / FN counts for the combined cut
  - Efficiency, Purity, Accuracy, F1
  - the ROC curve for "number of loud wires" at the current hit threshold,
    with a marker at the current wire-count threshold

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

    python interactive_hit_wire_count_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"]

categories = ("TP", "FN", "FP", "TN")
profiles_by_cat = {cat: [np.asarray(evt) for evt in wireProfile.get(cat, [])]
                    for cat in categories}

signal_profiles = profiles_by_cat["TP"] + profiles_by_cat["FN"]
background_profiles = profiles_by_cat["FP"] + profiles_by_cat["TN"]


def n_loud_wires(profiles, hit_thr):
    """Number of wires with hit count >= hit_thr, per event."""
    return np.array([np.sum(evt >= hit_thr) for evt in profiles], dtype=float)


# ============================================================
# Determine slider ranges
# ============================================================
all_profiles = signal_profiles + background_profiles

max_hits_in_wire = max(
    (evt.max() if evt.size > 0 else 0) for evt in all_profiles
)
hit_thr_min, hit_thr_max = 1, max(int(max_hits_in_wire), 1)

# Upper bound for wire-count slider: max possible "loud wires" at the most
# permissive hit threshold (hit_thr = 1), since raising hit_thr can only
# reduce the loud-wire count for a given event.
loose_signal_counts = n_loud_wires(signal_profiles, hit_thr_min)
loose_background_counts = n_loud_wires(background_profiles, hit_thr_min)
wire_count_max = max(loose_signal_counts.max(), loose_background_counts.max())
wire_count_min = 0.0


def metrics_at_thresholds(signal_counts, background_counts, wire_thr):
    """Classify as 'signal' when n_loud_wires >= wire_thr."""
    TP = np.sum(signal_counts >= wire_thr)
    FN = np.sum(signal_counts < wire_thr)
    FP = np.sum(background_counts >= wire_thr)
    TN = np.sum(background_counts < wire_thr)

    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=(12, 8))
gs = fig.add_gridspec(2, 2, height_ratios=[3, 1])
ax_hist = fig.add_subplot(gs[0, 0])
ax_roc = fig.add_subplot(gs[0, 1])
ax_text = fig.add_subplot(gs[1, :])
ax_text.axis("off")

plt.subplots_adjust(bottom=0.28)

hist_artists = {"signal": None, "background": None}
shade = {"patch": None}


def redraw_histogram(signal_counts, background_counts, wire_thr):
    ax_hist.clear()

    lo = 0
    hi = max(signal_counts.max(), background_counts.max(), 1)
    nbins = int(hi - lo + 1)
    edges = np.arange(lo, hi + 2) - 0.5  # integer-centered bins

    ax_hist.hist(signal_counts, bins=edges, density=True, alpha=0.5,
                 color="tab:green", label=f"Signal, n={len(signal_counts)}")
    ax_hist.hist(background_counts, bins=edges, density=True, alpha=0.5,
                 color="tab:red", label=f"Background, n={len(background_counts)}")

    ax_hist.axvline(wire_thr, color="k", ls="--", lw=2)
    ax_hist.axvspan(wire_thr, hi + 1, color="gray", alpha=0.15)

    ax_hist.set_xlabel("Number of loud wires (hits >= hit threshold)")
    ax_hist.set_ylabel("Probability density")
    ax_hist.set_title("Loud-wire count: signal vs background")
    ax_hist.grid(alpha=0.3)
    ax_hist.legend(loc="upper right", fontsize=8)


roc_line, = ax_roc.plot([], [], label="ROC (this hit threshold)")
diag_line, = ax_roc.plot([0, 1], [0, 1], "k--", alpha=0.6)
roc_point, = ax_roc.plot([], [], marker="*", markersize=18, color="red", zorder=5,
                          label="Current wire-count threshold")
ax_roc.set_xlabel("False Positive Rate")
ax_roc.set_ylabel("True Positive Rate")
ax_roc.set_title("ROC curve (recomputed per hit threshold)")
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_hit = plt.axes([0.15, 0.13, 0.7, 0.03])
ax_slider_wire = plt.axes([0.15, 0.06, 0.7, 0.03])

init_hit_thr = hit_thr_min
init_wire_thr = float(np.median(loose_signal_counts))

slider_hit = Slider(ax_slider_hit, "Min hits / wire", hit_thr_min, hit_thr_max,
                     valinit=init_hit_thr, valstep=1)
slider_wire = Slider(ax_slider_wire, "Min # loud wires", wire_count_min, wire_count_max,
                      valinit=init_wire_thr, valstep=1)

# Cache of counts for the current hit threshold (recomputed only when
# slider_hit moves, since it's the expensive part)
state = {"signal_counts": None, "background_counts": None, "hit_thr": None}


def recompute_counts(hit_thr):
    state["signal_counts"] = n_loud_wires(signal_profiles, hit_thr)
    state["background_counts"] = n_loud_wires(background_profiles, hit_thr)
    state["hit_thr"] = hit_thr


def update_common():
    wire_thr = slider_wire.val
    signal_counts = state["signal_counts"]
    background_counts = state["background_counts"]

    TP, FP, TN, FN, eff, pur, acc, F1, fpr_pt, tpr_pt = metrics_at_thresholds(
        signal_counts, background_counts, wire_thr
    )

    redraw_histogram(signal_counts, background_counts, wire_thr)

    y_true = np.concatenate([np.ones(len(signal_counts)), np.zeros(len(background_counts))])
    scores = np.concatenate([signal_counts, background_counts])
    fpr, tpr, _ = roc_curve(y_true, scores)
    auc = roc_auc_score(y_true, scores)
    roc_line.set_data(fpr, tpr)
    roc_line.set_label(f"ROC (AUC={auc:.3f})")
    roc_point.set_data([fpr_pt], [tpr_pt])
    ax_roc.legend(loc="lower right", fontsize=8)

    metrics_text.set_text(
        f"Min hits/wire = {int(state['hit_thr'])}    Min # loud wires = {int(wire_thr)}\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()


def on_hit_slider(val):
    recompute_counts(int(slider_hit.val))
    update_common()


def on_wire_slider(val):
    update_common()


slider_hit.on_changed(on_hit_slider)
slider_wire.on_changed(on_wire_slider)

recompute_counts(init_hit_thr)
update_common()

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 the number of
#   wires whose hit count is >= "Min hits / wire" is itself
#   >= "Min # loud wires".
# - Moving "Min hits / wire" changes what counts as a loud wire, so the
#   histogram, ROC curve, and AUC are all recomputed for that new
#   definition. Moving "Min # loud wires" only moves the operating point
#   along the (already computed) distribution/ROC for the current hit
#   threshold.