"""Pass 2: does training longer / regularising less (Guo et al. 2017), or top-N renormalisation,
produce the 'always sure' wall on CLINC150?  Same data, same TF-IDF config, seed 0.

Models (150 in-scope intents):
  orig  : pass-1 MLP (256 hidden, max_iter=40). Captured as the epoch-40 snapshot of the
          converged run below: identical seed/shuffles/optimizer path, verified against
          pass-1 gate_inputs/raw.json (in-scope test confidences must match).
  lr    : LogisticRegression(C=1000, max_iter=2000)  (effectively unregularised)
  conv  : MLPClassifier((256,), max_iter=MAX_ITER_CONV, n_iter_no_change=200, tol=0, random_state=0)
Logits = log(proba + 1e-12); temperature T fitted on in-scope val by NLL (bounded 0.05-10).
Run: python -I run_pass2.py   (optional: MAX_ITER_CONV env var, default 200)
"""
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.linear_model import LogisticRegression
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"]
MAX_ITER_CONV = int(os.environ.get("MAX_ITER_CONV", "200"))
SNAP_EPOCHS = sorted({10, 20, 40, 80, 120, 160, MAX_ITER_CONV} & set(range(1, MAX_ITER_CONV + 1)))

def pct(x): return round(100.0 * float(x), 1)
def log(*a): print(*a, flush=True)

data = json.loads((HERE / "data" / "data_full.json").read_text())
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"); Xote, _ = split("oos_test")
vec = TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True)
Xtr_m = vec.fit_transform(Xtr)
Xva_m, Xte_m, Xoos_m, Xpr_m = (vec.transform(x) for x in (Xva, Xte, Xote, PROBES))
log(f"TF-IDF features: {Xtr_m.shape[1]}")

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):
    L = logits(p)
    nll = lambda T: -log_softmax(L / T, axis=1)[np.arange(len(y)), y].mean()
    res = minimize_scalar(nll, bounds=(0.05, 10), method="bounded")
    return float(res.x)
def ece_bins(conf, correct, n_bins=15):
    edges = np.linspace(0, 1, n_bins + 1); ece = 0.0; rows = []
    for i in range(n_bins):
        m = (conf > edges[i]) & (conf <= edges[i + 1]) if i else (conf >= 0) & (conf <= edges[1])
        n = int(m.sum())
        if n: ece += n / len(conf) * abs(correct[m].mean() - conf[m].mean())
        rows.append(dict(bin=f"{edges[i]:.3f}-{edges[i+1]:.3f}", count=n,
                         mean_conf=round(float(conf[m].mean()), 3) if n else None,
                         accuracy=round(float(correct[m].mean()), 3) if n else None))
    return float(ece), rows
def top2(p): s = np.sort(p, axis=1); return s[:, -1], s[:, -1] - s[:, -2]
def route(p):
    conf, margin = top2(p)
    return np.where(conf < 0.40, "escalate", np.where((conf >= 0.80) & (margin >= 0.15), "act", "confirm"))
def split_pct(r): return {k: pct((r == k).mean()) for k in ("act", "confirm", "escalate")}
def train_xent(p_train, y_idx): return float(-np.log(p_train[np.arange(len(y_idx)), y_idx] + 1e-12).mean())

def evaluate(name, P_tr, P_va, P_te, P_oos, P_pr, classes, extra=None):
    idx = {c: i for i, c in enumerate(classes)}
    ytr_i, yva_i, yte_i = (np.array([idx[y] for y in ys]) for ys in (ytr, yva, yte))
    corr = P_te.argmax(1) == yte_i; conf = P_te.max(1); oconf = P_oos.max(1)
    ece, rows = ece_bins(conf, corr)
    top = conf > 14 / 15
    T = fit_T(P_va, yva_i)
    C_te, C_oos, C_pr = scale(P_te, T), scale(P_oos, T), scale(P_pr, T)
    ccorr = C_te.argmax(1) == yte_i
    ece_after, rows_after = ece_bins(C_te.max(1), ccorr)
    out = dict(
        inscope_acc_pct=pct(corr.mean()), train_xent=round(train_xent(P_tr, ytr_i), 4),
        inscope_ge_0_9_pct=pct((conf >= 0.9).mean()),
        oos_median=round(float(np.median(oconf)), 3), oos_ge_0_8_pct=pct((oconf >= 0.8).mean()),
        oos_ge_0_9_pct=pct((oconf >= 0.9).mean()), ece=round(ece, 4), bins=rows,
        bin_0_933_1=dict(count=int(top.sum()), mean_conf=round(float(conf[top].mean()), 3),
                         accuracy=round(float(corr[top].mean()), 3)),
        oos_top3=[[str(classes[i]), int(n)] for i, n in Counter(P_oos.argmax(1)).most_common(3)],
        T=round(T, 4), ece_after_T=round(ece_after, 4), acc_after_T_pct=pct(ccorr.mean()),
        inscope_ge_0_9_after_T_pct=pct((C_te.max(1) >= 0.9).mean()),
        oos_median_after_T=round(float(np.median(C_oos.max(1))), 3),
        oos_ge_0_8_after_T_pct=pct((C_oos.max(1) >= 0.8).mean()),
        oos_ge_0_9_after_T_pct=pct((C_oos.max(1) >= 0.9).mean()),
        probes_before=[dict(text=t, pred=str(classes[p.argmax()]), conf=round(float(p.max()), 3)) for t, p in zip(PROBES, P_pr)],
        probes_after_T=[dict(text=t, pred=str(classes[p.argmax()]), conf=round(float(p.max()), 3),
                             route=str(route(p[None])[0])) for t, p in zip(PROBES, C_pr)],
    )
    if extra: out.update(extra)
    _cache[name] = dict(P_te=P_te, P_oos=P_oos, P_pr=P_pr, C_te=C_te, C_oos=C_oos, C_pr=C_pr,
                        corr=corr, ccorr=ccorr, classes=classes)
    log(f"\n=== {name} ===")
    for k, v in out.items():
        if k not in ("bins",): log(f"  {k}: {v}")
    return out
_cache = {}

# ---------- converged MLP, with snapshots (epoch 40 == pass-1 model) ----------
class SnapMLP(MLPClassifier):
    snaps = {}
    def _update_no_improvement_count(self, early_stopping, X, y, sample_weight):
        if self.n_iter_ in SNAP_EPOCHS:
            SnapMLP.snaps[self.n_iter_] = dict(
                loss=float(self.loss_), P_tr=self.predict_proba(Xtr_m), P_va=self.predict_proba(Xva_m),
                P_te=self.predict_proba(Xte_m), P_oos=self.predict_proba(Xoos_m), P_pr=self.predict_proba(Xpr_m),
                t=round(time.time() - T0, 1))
            log(f"  [snapshot] epoch {self.n_iter_}: loss_={self.loss_:.6f}  t={time.time()-T0:.0f}s")
        return super()._update_no_improvement_count(early_stopping, X, y, sample_weight)

t = time.time()
conv = SnapMLP(hidden_layer_sizes=(256,), max_iter=MAX_ITER_CONV, n_iter_no_change=200, tol=0, random_state=0)
conv.fit(Xtr_m, ytr)
conv_sec = time.time() - t
log(f"Converged MLP: {conv.n_iter_} iterations, final loss_={conv.loss_:.6f}, {conv_sec:.1f}s")

# verify epoch-40 snapshot reproduces pass-1 model exactly
s40 = SnapMLP.snaps[40]
ref = json.load(open(HERE / "gate_inputs" / "raw.json"))
maxdiff = float(np.abs(s40["P_te"].max(1) - np.array(ref["conf"])).max())
log(f"Epoch-40 snapshot vs pass-1 in-scope confidences: max abs diff = {maxdiff:.2e}")
R["orig_reproduction_check"] = dict(max_abs_diff_vs_pass1_conf=maxdiff, reproduced=maxdiff < 1e-9)

cls = conv.classes_
R["orig_mlp_40"] = evaluate("orig MLP (epoch 40 = pass 1)", s40["P_tr"], s40["P_va"], s40["P_te"], s40["P_oos"], s40["P_pr"], cls,
                            dict(mlp_loss_=round(s40["loss"], 6), iterations=40))
final = SnapMLP.snaps[conv.n_iter_]
R["conv_mlp"] = evaluate(f"converged MLP (epoch {conv.n_iter_})", final["P_tr"], final["P_va"], final["P_te"], final["P_oos"], final["P_pr"], cls,
                         dict(mlp_loss_=round(float(conv.loss_), 6), iterations=int(conv.n_iter_), train_seconds=round(conv_sec, 1),
                              max_iter_used=MAX_ITER_CONV))

# epoch trajectory (Guo: NLL keeps falling after accuracy saturates)
idx = {c: i for i, c in enumerate(cls)}
yte_i = np.array([idx[y] for y in yte]); yva_i = np.array([idx[y] for y in yva]); ytr_i = np.array([idx[y] for y in ytr])
traj = []
for ep in SNAP_EPOCHS:
    s = SnapMLP.snaps[ep]
    corr = s["P_te"].argmax(1) == yte_i
    T = fit_T(s["P_va"], yva_i)
    traj.append(dict(epoch=ep, mlp_loss_=round(s["loss"], 6), train_acc_pct=pct((s["P_tr"].argmax(1) == ytr_i).mean()),
                     test_acc_pct=pct(corr.mean()),
                     val_nll=round(float(-np.log(s["P_va"][np.arange(len(yva_i)), yva_i] + 1e-12).mean()), 4),
                     ece=round(ece_bins(s["P_te"].max(1), corr)[0], 4), T=round(T, 4),
                     oos_median=round(float(np.median(s["P_oos"].max(1))), 3),
                     oos_ge_0_8_pct=pct((s["P_oos"].max(1) >= 0.8).mean()),
                     oos_ge_0_8_after_T_pct=pct((scale(s["P_oos"], T).max(1) >= 0.8).mean())))
R["mlp_epoch_trajectory"] = traj
log("\n=== MLP epoch trajectory ===")
for row in traj: log("  ", row)

# ---------- logistic regression, C=1000 ----------
t = time.time()
lr = LogisticRegression(C=1000, max_iter=2000)
lr.fit(Xtr_m, ytr)
lr_sec = time.time() - t
log(f"\nLogReg C=1000: n_iter={lr.n_iter_.tolist()}, {lr_sec:.1f}s")
assert list(lr.classes_) == list(cls)
R["lr_C1000"] = evaluate("LogReg C=1000", *(lr.predict_proba(m) for m in (Xtr_m, Xva_m, Xte_m, Xoos_m, Xpr_m)), lr.classes_,
                         dict(iterations=int(np.max(lr.n_iter_)), train_seconds=round(lr_sec, 1)))

# ---------- comparison table ----------
comp = []
for key, label in (("orig_mlp_40", "MLP 40 epochs (pass 1)"), ("lr_C1000", "LogReg C=1000"),
                   ("conv_mlp", f"MLP {R['conv_mlp']['iterations']} epochs, tol=0")):
    m = R[key]
    comp.append(dict(model=label, inscope_acc_pct=m["inscope_acc_pct"], train_xent=m["train_xent"],
                     mlp_loss_=m.get("mlp_loss_"), ece=m["ece"], T=m["T"], ece_after_T=m["ece_after_T"],
                     oos_ge_0_8_raw_pct=m["oos_ge_0_8_pct"], oos_ge_0_8_after_T_pct=m["oos_ge_0_8_after_T_pct"],
                     oos_median=m["oos_median"]))
R["step1_comparison"] = comp
log("\n=== Step 1 comparison ===")
for row in comp: log("  ", row)

# ---------- step 2: top-N renormalisation on pass-1 MLP ----------
def renorm_topk(p, k):
    out = np.zeros_like(p); idx_k = np.argsort(-p, axis=1)[:, :k]
    rows = np.arange(len(p))[:, None]
    out[rows, idx_k] = p[rows, idx_k]
    return out / out.sum(1, keepdims=True)
o = _cache["orig MLP (epoch 40 = pass 1)"]
renorm = {}
for label, f in (("none", lambda p: p), ("top3", lambda p: renorm_topk(p, 3)), ("top2", lambda p: renorm_topk(p, 2))):
    pi, po = f(o["P_te"]), f(o["P_oos"])
    ci, co = pi.max(1), po.max(1)
    assert (pi.argmax(1) == o["P_te"].argmax(1)).all()
    renorm[label] = dict(inscope_ge_0_8_pct=pct((ci >= 0.8).mean()), inscope_ge_0_9_pct=pct((ci >= 0.9).mean()),
                         oos_ge_0_8_pct=pct((co >= 0.8).mean()), oos_ge_0_9_pct=pct((co >= 0.9).mean()),
                         oos_median=round(float(np.median(co)), 3), oos_mean=round(float(co.mean()), 3),
                         inscope_ece=round(ece_bins(ci, o["corr"])[0], 4),
                         probes=[dict(text=t, conf=round(float(p.max()), 3)) for t, p in zip(PROBES, f(o["P_pr"]))])
R["step2_topN_renorm"] = renorm
log("\n=== Step 2: top-N renormalisation (pass-1 MLP) ===")
for k, v in renorm.items(): log(f"  {k}: {v}")

# ---------- step 3: margin routing on model with T closest to 1 ----------
best_key = min(("lr_C1000", "conv_mlp"), key=lambda k: abs(R[k]["T"] - 1))
best_name = {"lr_C1000": "LogReg C=1000", "conv_mlp": f"converged MLP (epoch {conv.n_iter_})"}[best_key]
b = _cache[best_name]
r_oos, r_in = route(b["C_oos"]), route(b["C_te"])
R["step3_margin_routing"] = dict(model=best_key, T=R[best_key]["T"], oos_routing_pct=split_pct(r_oos),
                                 inscope_routing_pct=split_pct(r_in),
                                 inscope_acc_among_act_pct=pct(b["ccorr"][r_in == "act"].mean()))
log("\n=== Step 3: margin routing ===", R["step3_margin_routing"])
# also for the other step-1 model, for context
other = "conv_mlp" if best_key == "lr_C1000" else "lr_C1000"
on = {"lr_C1000": "LogReg C=1000", "conv_mlp": f"converged MLP (epoch {conv.n_iter_})"}[other]
c = _cache[on]; ro, ri = route(c["C_oos"]), route(c["C_te"])
R["step3_margin_routing_other_model"] = dict(model=other, T=R[other]["T"], oos_routing_pct=split_pct(ro),
                                             inscope_routing_pct=split_pct(ri),
                                             inscope_acc_among_act_pct=pct(c["ccorr"][ri == "act"].mean()))
log("  (other model)", R["step3_margin_routing_other_model"])

# ---------- step 4: gate on converged MLP raw and calibrated ----------
gate_dir = HERE / "gate_inputs"
cm = _cache[f"converged MLP (epoch {conv.n_iter_})"]
def dump(name, P_te, corr, P_pr):
    json.dump(dict(conf=[float(x) for x in P_te.max(1)], correct=[bool(x) for x in corr],
                   probes=[dict(text=t, pred=str(cls[p.argmax()]), conf=float(p.max())) for t, p in zip(PROBES, P_pr)]),
              open(gate_dir / f"{name}.json", "w"))
dump("conv_raw", cm["P_te"], cm["corr"], cm["P_pr"])
dump("conv_calibrated", cm["C_te"], cm["ccorr"], cm["C_pr"])
G = {}
for name in ("conv_raw", "conv_calibrated"):
    env = dict(os.environ, GATE_INPUT=str(gate_dir / f"{name}.json"))
    out = subprocess.run([sys.executable, "-I", "-m", "pytest", "-q", "-p", "no:cacheprovider", "-rf",
                          str(HERE / "test_confidence_gate.py")], env=env, capture_output=True, text=True, cwd=HERE)
    (HERE / f"pytest_pass2_{name}.log").write_text(out.stdout + out.stderr)
    lines = [l for l in out.stdout.splitlines() if l.strip()]
    last = lines[-1]
    G[name] = dict(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),
                   failing_tests=[l.split("::", 1)[1].split(" - ")[0] for l in lines if l.startswith("FAILED")])
G["conv_calibrated_probe_routes"] = [dict(text=t, conf=round(float(p.max()), 3),
                                          margin=round(float(top2(p[None])[1][0]), 3), route=str(route(p[None])[0]))
                                     for t, p in zip(PROBES, cm["C_pr"])]
G["conv_raw_ece"] = round(ece_bins(cm["P_te"].max(1), cm["corr"])[0], 4)
G["conv_calibrated_ece"] = round(ece_bins(cm["C_te"].max(1), cm["ccorr"])[0], 4)
R["step4_gate"] = G
log("\n=== Step 4: gate ===")
for k, v in G.items(): log(f"  {k}: {v}")

R["versions"] = dict(python=sys.version.split()[0], numpy=np.__version__, scipy=scipy.__version__, sklearn=sklearn.__version__)
R["runtime_seconds"] = round(time.time() - T0, 1)
(HERE / "results_pass2.json").write_text(json.dumps(R, indent=2))
log(f"\nTotal runtime: {R['runtime_seconds']}s")