"""Overconfident intent classifier experiment on CLINC150 (Larson et al., EMNLP 2019).

TF-IDF (word 1-2 grams, sublinear_tf) + MLPClassifier(256 hidden, random_state=0).
Confidence = predict_proba top-1. Logits = log(proba + 1e-12); since softmax(log p) = p,
softmax(log p / T) is exactly temperature scaling applied to the model's probabilities.
All randomness fixed (random_state=0, np seed 0). Run: python -I run_experiment.py
"""
import json, os, subprocess, sys, time, warnings
from collections import Counter
from pathlib import Path

import numpy as np
import scipy, sklearn
from scipy.optimize import minimize_scalar
from scipy.special import softmax, log_softmax
from sklearn.exceptions import ConvergenceWarning
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.neural_network import MLPClassifier

warnings.filterwarnings("ignore", category=ConvergenceWarning)
np.random.seed(0)
HERE = Path(__file__).resolve().parent
T0 = time.time()
R = {}

PROBES = ["what's the capital of France?", "asdf jkl;", "?",
          "write me a poem about mortgages", "can you help me with my thing"]
THRESHOLDS = [round(x, 2) for x in np.arange(0.50, 0.951, 0.05)]
HIDDEN, MAX_ITER = 256, 40

def pct(x): return round(100.0 * float(x), 1)

# ---------------- data ----------------
data = json.loads((HERE / "data" / "data_full.json").read_text())
counts = {k: len(v) for k, v in data.items()}
R["dataset_counts"] = counts
print("Dataset counts:", counts)
expected = dict(train=15000, val=3000, test=4500, oos_train=100, oos_val=100, oos_test=1000)
for k, v in expected.items():
    assert counts[k] == v, (k, counts[k], v)
def split(key): return [t for t, _ in data[key]], [l for _, l in data[key]]
Xtr, ytr = split("train"); Xva, yva = split("val"); Xte, yte = split("test")
Xotr, _ = split("oos_train"); Xova, _ = split("oos_val"); Xote, _ = split("oos_test")
assert len(set(ytr)) == 150

# ---------------- helpers ----------------
def train(texts, labels):
    vec = TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True)
    X = vec.fit_transform(texts)
    clf = MLPClassifier(hidden_layer_sizes=(HIDDEN,), max_iter=MAX_ITER, random_state=0)
    t = time.time(); clf.fit(X, labels)
    return vec, clf, X.shape[1], time.time() - t

def proba(vec, clf, texts): return clf.predict_proba(vec.transform(texts))
def logits(p): return np.log(p + 1e-12)
def scale(p, T): return softmax(logits(p) / T, axis=1)

def fit_T(p, y_idx):
    L = logits(p)
    nll = lambda T: -log_softmax(L / T, axis=1)[np.arange(len(y_idx)), y_idx].mean()
    res = minimize_scalar(nll, bounds=(0.05, 10), method="bounded")
    return float(res.x), float(nll(1.0)), float(res.fun)

def reliability(conf, correct, n_bins=15):
    edges = np.linspace(0, 1, n_bins + 1); rows = []; ece = 0.0
    for i in range(n_bins):
        lo, hi = edges[i], edges[i + 1]
        m = (conf > lo) & (conf <= hi) if i > 0 else (conf >= lo) & (conf <= hi)
        n = int(m.sum())
        if n:
            mc, acc = float(conf[m].mean()), float(correct[m].mean())
            ece += n / len(conf) * abs(acc - mc)
        else:
            mc = acc = None
        rows.append(dict(bin=f"{lo:.3f}-{hi:.3f}", count=n, mean_conf=mc, accuracy=acc))
    return rows, float(ece)

def sweep(conf_in, correct_in, conf_oos):
    out = []
    for t in THRESHOLDS:
        out.append(dict(threshold=t, oos_recall_pct=pct((conf_oos < t).mean()),
                        inscope_false_escalation_pct=pct((conf_in[correct_in] < t).mean())))
    first = next((r for r in out if r["oos_recall_pct"] >= 50.0), None)
    return out, first

def top2(p):
    s = np.sort(p, axis=1); return s[:, -1], s[:, -1] - s[:, -2]

def route(p, pred_labels=None):
    conf, margin = top2(p)
    r = np.where(conf < 0.40, "escalate",
                 np.where((conf >= 0.80) & (margin >= 0.15), "act", "confirm"))
    if pred_labels is not None:  # explicit oos class prediction always escalates
        r = np.where(np.asarray(pred_labels) == "oos", "escalate", r)
    return r

def split_pct(r): return {k: pct((r == k).mean()) for k in ("act", "confirm", "escalate")}

def print_table(rows, cols):
    print("  " + " | ".join(cols))
    for row in rows:
        print("  " + " | ".join(
            (f"{row[c]:.3f}" if isinstance(row[c], float) else str(row[c])) for c in cols))

# ---------------- model 1: 150 in-scope intents ----------------
vec1, clf1, nfeat1, tsec1 = train(Xtr, ytr)
R["model1"] = dict(n_features=nfeat1, train_seconds=round(tsec1, 1), n_iter=int(clf1.n_iter_))
print(f"\nModel 1 trained: {nfeat1} features, {clf1.n_iter_} epochs, {tsec1:.1f}s")
cls1 = clf1.classes_; idx1 = {c: i for i, c in enumerate(cls1)}
P_te, P_va, P_oos = proba(vec1, clf1, Xte), proba(vec1, clf1, Xva), proba(vec1, clf1, Xote)
P_probe = proba(vec1, clf1, PROBES)
yte_i = np.array([idx1[y] for y in yte]); yva_i = np.array([idx1[y] for y in yva])
pred_te = P_te.argmax(1); correct_te = pred_te == yte_i
conf_te = P_te.max(1); conf_oos = P_oos.max(1)

# A
A = dict(
    inscope_test_accuracy_pct=pct(correct_te.mean()),
    inscope_share_ge_0_9_pct=pct((conf_te >= 0.9).mean()),
    oos_median_conf=round(float(np.median(conf_oos)), 3),
    oos_mean_conf=round(float(conf_oos.mean()), 3),
    oos_share_ge_0_8_pct=pct((conf_oos >= 0.8).mean()),
    oos_share_ge_0_9_pct=pct((conf_oos >= 0.9).mean()),
    oos_top3_landing_intents=[[cls1[i], n] for i, n in Counter(P_oos.argmax(1)).most_common(3)],
    probes=[dict(text=t, pred=cls1[p.argmax()], conf=round(float(p.max()), 3)) for t, p in zip(PROBES, P_probe)],
)
R["A_raw"] = A
print("\n=== A. Raw model: the 'always sure' wall ===")
for k, v in A.items():
    if k != "probes": print(f"  {k}: {v}")
for pr in A["probes"]: print(f"  probe {pr['text']!r} -> {pr['pred']} @ {pr['conf']}")

# B
rows_b, ece_b = reliability(conf_te, correct_te)
top_bin = dict(count=int((conf_te >= 0.9).sum()), accuracy_pct=pct(correct_te[conf_te >= 0.9].mean()),
               mean_conf=round(float(conf_te[conf_te >= 0.9].mean()), 3))
R["B_calibration_before"] = dict(ece=round(ece_b, 4), bins=rows_b, bin_0_9_to_1_0=top_bin)
print(f"\n=== B. Calibration before (in-scope test, 15 bins) === ECE={ece_b:.4f}")
print_table(rows_b, ["bin", "count", "mean_conf", "accuracy"])
print("  0.9-1.0 range:", top_bin)

# C
sw_raw, first_raw = sweep(conf_te, correct_te, conf_oos)
R["C_threshold_raw"] = dict(sweep=sw_raw, first_ge_50pct_oos_recall=first_raw)
print("\n=== C. Threshold alone (raw scores) ===")
print_table(sw_raw, ["threshold", "oos_recall_pct", "inscope_false_escalation_pct"])
print("  first >=50% OOS recall:", first_raw)

# D
T1, nll1_before, nll1_after = fit_T(P_va, yva_i)
Pc_te, Pc_oos, Pc_probe = scale(P_te, T1), scale(P_oos, T1), scale(P_probe, T1)
cconf_te, cconf_oos = Pc_te.max(1), Pc_oos.max(1)
ccorrect_te = Pc_te.argmax(1) == yte_i
rows_d, ece_d = reliability(cconf_te, ccorrect_te)
sw_cal, first_cal = sweep(cconf_te, ccorrect_te, cconf_oos)
D = dict(T=round(T1, 4), val_nll_before=round(nll1_before, 4), val_nll_after=round(nll1_after, 4),
         ece_after=round(ece_d, 4), accuracy_after_pct=pct(ccorrect_te.mean()),
         accuracy_identical=bool((Pc_te.argmax(1) == pred_te).all()),
         inscope_share_ge_0_9_after_pct=pct((cconf_te >= 0.9).mean()),
         oos_median_conf_after=round(float(np.median(cconf_oos)), 3),
         oos_share_ge_0_8_after_pct=pct((cconf_oos >= 0.8).mean()),
         oos_share_ge_0_9_after_pct=pct((cconf_oos >= 0.9).mean()),
         bins_after=rows_d, sweep=sw_cal, first_ge_50pct_oos_recall=first_cal,
         probes_after=[dict(text=t, pred=cls1[p.argmax()], conf=round(float(p.max()), 3)) for t, p in zip(PROBES, Pc_probe)])
R["D_temperature"] = D
print("\n=== D. Temperature scaling ===")
for k, v in D.items():
    if k not in ("bins_after", "sweep", "probes_after"): print(f"  {k}: {v}")
print_table(rows_d, ["bin", "count", "mean_conf", "accuracy"])
print_table(sw_cal, ["threshold", "oos_recall_pct", "inscope_false_escalation_pct"])
for pr in D["probes_after"]: print(f"  probe {pr['text']!r} -> {pr['pred']} @ {pr['conf']}")

# E(i)
r_oos, r_in = route(Pc_oos), route(Pc_te)
act_in = r_in == "act"
Ei = dict(oos_routing_pct=split_pct(r_oos), inscope_routing_pct=split_pct(r_in),
          inscope_accuracy_among_act_pct=pct(ccorrect_te[act_in].mean()))
R["E_i_margin_routing_calibrated"] = Ei
print("\n=== E(i). Margin routing on calibrated 150-class scores ===")
for k, v in Ei.items(): print(f"  {k}: {v}")

# E(ii): 151-class model
vec2, clf2, nfeat2, tsec2 = train(Xtr + Xotr, ytr + ["oos"] * len(Xotr))
cls2 = clf2.classes_; idx2 = {c: i for i, c in enumerate(cls2)}
print(f"\nModel 2 (151 classes) trained: {nfeat2} features, {clf2.n_iter_} epochs, {tsec2:.1f}s")
Q_va = proba(vec2, clf2, Xva + Xova); yq_va = np.array([idx2[y] for y in yva + ["oos"] * len(Xova)])
T2, nll2_before, nll2_after = fit_T(Q_va, yq_va)
Q_te, Q_oos, Q_probe = (scale(proba(vec2, clf2, X), T2) for X in (Xte, Xote, PROBES))
Qraw_te, Qraw_oos = proba(vec2, clf2, Xte), proba(vec2, clf2, Xote)
yte2 = np.array([idx2[y] for y in yte])
qpred_te, qpred_oos = cls2[Q_te.argmax(1)], cls2[Q_oos.argmax(1)]
qcorrect_te = Q_te.argmax(1) == yte2
rq_oos, rq_in = route(Q_oos, qpred_oos), route(Q_te, qpred_te)
rows_q, ece_q = reliability(Q_te.max(1), qcorrect_te)
Eii = dict(T=round(T2, 4), val_nll_before=round(nll2_before, 4), val_nll_after=round(nll2_after, 4),
           oos_recall_pred_oos_pct=pct((qpred_oos == "oos").mean()),
           inscope_test_accuracy_pct=pct(qcorrect_te.mean()),
           inscope_wrongly_pred_oos_pct=pct((qpred_te == "oos").mean()),
           ece_inscope_test=round(ece_q, 4),
           raw_uncalibrated_ece_inscope_test=round(reliability(Qraw_te.max(1), Qraw_te.argmax(1) == yte2)[1], 4),
           oos_routing_pct=split_pct(rq_oos), inscope_routing_pct=split_pct(rq_in),
           inscope_accuracy_among_act_pct=pct(qcorrect_te[rq_in == "act"].mean()),
           probes=[dict(text=t, pred=cls2[p.argmax()], conf=round(float(p.max()), 3)) for t, p in zip(PROBES, Q_probe)])
R["E_ii_oos_class_calibrated"] = Eii
print("=== E(ii). 151-class model (oos class) + temperature ===")
for k, v in Eii.items():
    if k != "probes": print(f"  {k}: {v}")
for pr in Eii["probes"]: print(f"  probe {pr['text']!r} -> {pr['pred']} @ {pr['conf']}")

# E(iii)
qconf_oos = Q_oos.max(1)
miss = np.where((qpred_oos != "oos") & (qconf_oos >= 0.8))[0]
miss = miss[np.argsort(-qconf_oos[miss], kind="stable")]
Eiii = dict(n_oos_test_confident_misroutes=int(len(miss)),
            n_oos_test_misrouted_any_conf=int((qpred_oos != "oos").sum()),
            examples=[dict(text=Xote[i], intent=qpred_oos[i], conf=round(float(qconf_oos[i]), 3)) for i in miss[:5]])
R["E_iii_dead_end_2"] = Eiii
print("\n=== E(iii). What the 100-example oos class still misses (calibrated conf >= 0.8) ===")
print(f"  count: {Eiii['n_oos_test_confident_misroutes']} of 1000 (any-confidence misroutes: {Eiii['n_oos_test_misrouted_any_conf']})")
for ex in Eiii["examples"]: print(f"  {ex['text']!r} -> {ex['intent']} @ {ex['conf']}")

# F: CI gate
gate_dir = HERE / "gate_inputs"; gate_dir.mkdir(exist_ok=True)
def dump(name, conf, correct, probe_p, probe_cls):
    json.dump(dict(conf=[float(c) for c in conf], correct=[bool(c) for c in correct],
                   probes=[dict(text=t, pred=str(probe_cls[p.argmax()]), conf=float(p.max())) for t, p in zip(PROBES, probe_p)]),
              open(gate_dir / f"{name}.json", "w"))
dump("raw", conf_te, correct_te, P_probe, cls1)
dump("fixed", Q_te.max(1), qcorrect_te, Q_probe, cls2)
F = {}
for name in ("raw", "fixed"):
    env = dict(os.environ, GATE_INPUT=str(gate_dir / f"{name}.json"))
    out = subprocess.run([sys.executable, "-I", "-m", "pytest", "-q", "-p", "no:cacheprovider",
                          str(HERE / "test_confidence_gate.py")], env=env, capture_output=True, text=True, cwd=HERE)
    last = [l for l in out.stdout.strip().splitlines() if l.strip()][-1]
    passed = sum(int(w.split()[0]) for w in last.split(",") if "passed" in w)
    failed = sum(int(w.split()[0]) for w in last.split(",") if "failed" in w)
    F[name] = dict(passed=passed, failed=failed, summary=last.strip("= ").strip())
    (HERE / f"pytest_{name}.log").write_text(out.stdout + out.stderr)
R["F_ci_gate"] = F
print("\n=== F. CI gate (pytest, 6 tests: 1 ECE + 5 probes) ===")
for k, v in F.items(): print(f"  {k}: {v}")

R["versions"] = dict(python=sys.version.split()[0], numpy=np.__version__, scipy=scipy.__version__,
                     sklearn=sklearn.__version__)
R["config"] = dict(tfidf="word ngram (1,2), sublinear_tf=True, defaults otherwise", mlp_hidden=HIDDEN,
                   mlp_max_iter=MAX_ITER, random_state=0, logits="log(proba+1e-12)",
                   ece_bins=15, routing="escalate conf<0.40; act conf>=0.80 & margin>=0.15; else confirm; predicted oos -> escalate")
R["runtime_seconds"] = round(time.time() - T0, 1)
(HERE / "results.json").write_text(json.dumps(R, indent=2))
print(f"\nTotal runtime: {R['runtime_seconds']}s | versions: {R['versions']}")