"""
Confidence-based rejection (selective prediction) on the wild test.

The calibration analysis showed the studio model's wrong cross-source predictions
sit at low confidence. Here we quantify the practical consequence: if we abstain
on low-confidence predictions, accuracy on the accepted subset rises. This is an
honest deployment knob (it trades coverage for accuracy), not a fix for the
collapse. Reads calibration_s*.json (studio-trained model, wild test).
"""
import sys, os, glob, json
from pathlib import Path

_base = "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet"
sys.path.insert(0, _base)

import numpy as np
import figstyle

ROOT = Path(_base)
OUT = ROOT / "outputs" / "rejection"


def main():
    files = [f for f in sorted(glob.glob(str(ROOT / "outputs" / "calibration_s*.json")))
             if "smoke" not in Path(f).name.lower()]
    if not files:
        raise SystemExit("no calibration_s*.json — run eval_calibration.py first")
    conf, corr = [], []
    for f in files:
        d = json.load(open(f))["in"]["ss"]          # studio-trained on wild test
        conf += d["confidence"]; corr += d["correct"]
    conf, corr = np.asarray(conf), np.asarray(corr, float)

    order = np.argsort(-conf)                         # most confident first
    c = corr[order]
    k = np.arange(1, len(c) + 1)
    coverage = k / len(c)
    accuracy = np.cumsum(c) / k                       # accuracy of the accepted (top-conf) subset

    # coverage at which accepted-accuracy first reaches some targets
    marks = {}
    for tgt in (0.90, 0.95, 0.99):
        idx = np.where(accuracy >= tgt)[0]
        marks[tgt] = float(coverage[idx[-1]]) if len(idx) else 0.0

    json.dump({"n": int(len(c)), "base_accuracy": float(corr.mean()),
               "coverage_at_acc": {str(k): v for k, v in marks.items()}},
              open(ROOT / "outputs" / "rejection.json", "w"), indent=2)

    figstyle.apply()
    import matplotlib.pyplot as plt
    P = figstyle.PALETTE
    fig, ax = plt.subplots(figsize=(6.4, 4.4))
    ax.plot(coverage * 100, accuracy * 100, color=P["wild"], lw=2)
    ax.axhline(corr.mean() * 100, color=P["chance"], ls=":", lw=1,
               label=f"full-coverage acc {corr.mean()*100:.0f}%")
    for tgt, cov in marks.items():
        if cov > 0:
            ax.plot(cov * 100, tgt * 100, "o", color=P["cross_broken"], ms=6)
            ax.annotate(f"{tgt*100:.0f}% acc @ {cov*100:.0f}% coverage",
                        (cov * 100, tgt * 100), textcoords="offset points",
                        xytext=(6, -10), fontsize=8)
    ax.set_xlabel("coverage (% of wild predictions kept)")
    ax.set_ylabel("accuracy on accepted (%)")
    ax.set_xlim(0, 100); ax.set_ylim(50, 101)
    ax.set_title("Selective prediction on the wild test (studio-trained model)", fontsize=11)
    ax.legend(loc="lower left", fontsize=9); ax.grid(alpha=0.25)
    fig.tight_layout()
    figstyle.save(fig, str(OUT))
    print(f"base acc {corr.mean()*100:.1f}%  |  "
          + "  ".join(f">= {int(t*100)}% acc @ {v*100:.0f}% cov" for t, v in marks.items()))


if __name__ == "__main__":
    main()