"""
Geometric picture of the ROLL TPR@FPR objective.

Upper panel: score PDFs for both classes; shaded area under the positive-class
PDF to the left of τ equals the ROLL loss F̂₁(τ).
Lower panel: CDFs showing how τ is read from F̂₀ and how the loss and TPR
are read from F̂₁.

Run from the thesis root:
    nix develop --command python3 content/method/figures/roll_principle.py
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.stats import norm

OUT_DIR = "content/method/figures"

np.random.seed(3)

# --- Parameters ------------------------------------------------------------
mu0, sig0 = 0.0, 1.00    # negative class N(0, 1)
mu1, sig1 = 1.5, 1.00    # positive class N(1.5, 1)  — intentional overlap
alpha = 0.20             # target FPR

# τ = (1-α)-quantile of the negative class  →  F̂₀(τ) = 1-α  →  FPR = α
tau = norm.ppf(1 - alpha, mu0, sig0)
F0_tau = norm.cdf(tau, mu0, sig0)   # = 1 - alpha
F1_tau = norm.cdf(tau, mu1, sig1)   # ROLL loss
tpr = 1.0 - F1_tau

# --- Sample points for positive class (rug) --------------------------------
rug = np.random.normal(mu1, sig1, 35)

# --- Score axis ------------------------------------------------------------
x = np.linspace(-3.6, 5.0, 900)

# --- Palette ---------------------------------------------------------------
C_NEG  = "#4477AA"
C_POS  = "#CC6633"
C_TAU  = "#333333"
C_LOSS = "#7799BB"

# ── Figure ─────────────────────────────────────────────────────────────────
fig, (ax1, ax2) = plt.subplots(
    2, 1, figsize=(5.5, 4.8), sharex=True,
    gridspec_kw={"height_ratios": [1.25, 1.0]},
)
fig.subplots_adjust(hspace=0.07)

# ══ Top: PDF ══════════════════════════════════════════════════════════════

pdf0 = norm.pdf(x, mu0, sig0)
pdf1 = norm.pdf(x, mu1, sig1)

ax1.plot(x, pdf0, color=C_NEG, linewidth=1.8, label=r"Negative class ($y=0$)")
ax1.plot(x, pdf1, color=C_POS, linewidth=1.8, label=r"Positive class ($y=1$)")
ax1.axvline(tau, color=C_TAU, linestyle="--", linewidth=1.3,
            label=r"Threshold $\tau$")

# Shade: area under positive PDF to the left of τ = loss = F̂₁(τ)
x_left = x[x <= tau]
ax1.fill_between(x_left, norm.pdf(x_left, mu1, sig1),
                 color=C_LOSS, alpha=0.30, label=r"Loss $= \hat{F}_1(\tau)$")

# Rug plot
rug_y = np.full_like(rug, -0.006)
ax1.scatter(rug, rug_y, color=C_POS, s=18, marker="|",
            zorder=3, alpha=0.65, clip_on=False)

ax1.set_ylabel("PDF", fontsize=9)
ax1.set_yticks([])
ax1.set_ylim(bottom=-0.025)
ax1.legend(fontsize=7.5, loc="upper right", framealpha=0.88)

# ══ Bottom: CDF ═══════════════════════════════════════════════════════════

cdf0 = norm.cdf(x, mu0, sig0)
cdf1 = norm.cdf(x, mu1, sig1)

ax2.plot(x, cdf0, color=C_NEG, linewidth=1.8, label=r"$\hat{F}_0$ (neg.)")
ax2.plot(x, cdf1, color=C_POS, linewidth=1.8, label=r"$\hat{F}_1$ (pos.)")
ax2.axvline(tau, color=C_TAU, linestyle="--", linewidth=1.3)

# Horizontal reference lines
ax2.axhline(F0_tau, color=C_NEG, linestyle="-.", linewidth=0.9, alpha=0.7)
ax2.axhline(F1_tau, color=C_POS, linestyle="-.", linewidth=0.9, alpha=0.7)

# Dots at the intercepts
ax2.plot(tau, F0_tau, "o", color=C_NEG, markersize=5, zorder=4)
ax2.plot(tau, F1_tau, "o", color=C_POS, markersize=5, zorder=4)

# Right-side labels for the intercepts (placed where CDFs have flattened to ~1)
x_label_right = x[-1] - 0.2
ax2.text(x_label_right, F0_tau, rf"$1-\alpha$",
         fontsize=8, color=C_NEG, va="bottom", ha="right")
ax2.text(x_label_right, F1_tau, r"$\hat{F}_1(\tau)$" + "\n(loss)",
         fontsize=7.5, color=C_POS, va="top", ha="right")

# TPR annotation: double-headed arrow on the right + label
x_bracket = tau + 1.9
ax2.annotate("", xy=(x_bracket, 1.0), xytext=(x_bracket, F1_tau),
             arrowprops=dict(arrowstyle="<->", color="black",
                             lw=0.8, mutation_scale=10))
ax2.text(x_bracket + 0.12, (1.0 + F1_tau) / 2,
         f"TPR\n= {tpr*100:.0f}%",
         fontsize=7.5, va="center", ha="left", color="black")

ax2.set_ylabel("CDF", fontsize=9)
ax2.set_xlabel(r"Score $f_\theta(\mathbf{x})$", fontsize=9)
ax2.set_yticks([0.0, 0.25, 0.50, 0.75, 1.0])
ax2.set_yticklabels(["0", "0.25", "0.5", "0.75", "1"], fontsize=7.5)
ax2.set_xlim(x[0] - 0.3, x[-1])
ax2.legend(fontsize=7.5, loc="upper left", framealpha=0.88)

for ax in (ax1, ax2):
    ax.tick_params(labelsize=8)
    ax.spines[["top", "right"]].set_visible(False)

fig.savefig(f"{OUT_DIR}/roll_principle.pdf", bbox_inches="tight")
fig.savefig(f"{OUT_DIR}/roll_principle.png", bbox_inches="tight", dpi=180)
print(f"Saved to {OUT_DIR}/roll_principle.{{pdf,png}}")
print(f"tau={tau:.3f}  F0(tau)={F0_tau:.3f}  F1(tau)={F1_tau:.3f}  TPR={tpr:.3f}")