"""
Phase 3 — Ensemble, calibrated thresholding, advanced statistics.

Inputs: per-model *_test_preds.json (containing labels/preds/probs) for every
model trained in v2 and foundation. Produces:
  * ensemble (soft-vote over selected models + weighted by val-F1)
  * per-class threshold optimization on val (maximize macro-F1)
  * conformal prediction sets at 90% coverage (Mondrian by class)
  * Bonferroni-corrected pairwise McNemar
  * per-class bootstrap CIs
  * Cohen's kappa, Brier score
"""
import argparse, json, glob, os
from pathlib import Path
import numpy as np
from scipy.stats import binom
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support,
    roc_auc_score, average_precision_score, cohen_kappa_score, brier_score_loss
)


def load_all(results_dir):
    preds = {}
    for f in sorted(glob.glob(os.path.join(results_dir, "*_test_preds.json"))):
        name = Path(f).stem.replace("_test_preds", "")
        d = json.load(open(f))
        preds[name] = {
            "labels": np.array(d["labels"]),
            "preds": np.array(d["preds"]),
            "probs": np.array(d["probs"]),
        }
    return preds


def ece(probs, labels, n_bins=15):
    conf = probs.max(1); pred = probs.argmax(1); correct = (pred == labels).astype(float)
    bins = np.linspace(0, 1, n_bins+1); e = 0.0
    for i in range(n_bins):
        m = (conf > bins[i]) & (conf <= bins[i+1])
        if m.sum(): e += m.mean() * abs(correct[m].mean() - conf[m].mean())
    return float(e)


def macro_brier(probs, labels, K):
    oh = np.zeros_like(probs); oh[np.arange(len(labels)), labels] = 1
    return float(((probs - oh) ** 2).sum(1).mean())


def ensemble_soft(preds_dict, weights=None):
    names = list(preds_dict.keys())
    if weights is None: weights = {n: 1.0 for n in names}
    labels = preds_dict[names[0]]["labels"]
    probs = np.zeros_like(preds_dict[names[0]]["probs"])
    wsum = 0
    for n in names:
        probs += weights[n] * preds_dict[n]["probs"]; wsum += weights[n]
    probs /= wsum
    return labels, probs


def per_class_thresholds(val_probs, val_labels, K, n_thr=51):
    """Find scalar bias per class that maximizes macro-F1 on val."""
    thresholds = np.linspace(-0.3, 0.3, n_thr)
    best = np.zeros(K)
    for k in range(K):
        best_f1, best_t = -1, 0.0
        for t in thresholds:
            biased = val_probs.copy(); biased[:, k] += t
            preds = biased.argmax(1)
            _, _, f1, _ = precision_recall_fscore_support(val_labels, preds, average="macro", zero_division=0)
            if f1 > best_f1: best_f1, best_t = f1, t
        best[k] = best_t
    return best


def apply_thresholds(probs, biases):
    biased = probs.copy() + biases[None, :]
    return biased.argmax(1)


def bootstrap_ci_acc(labels, preds, n=2000, seed=42):
    rng = np.random.default_rng(seed); N = len(labels); vals = []
    for _ in range(n):
        idx = rng.integers(0, N, N)
        vals.append((preds[idx] == labels[idx]).mean())
    return float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))


def mcnemar_pair(labels, p1, p2):
    c1 = p1 == labels; c2 = p2 == labels
    b = int((c1 & ~c2).sum()); c = int((~c1 & c2).sum()); n = b + c
    if n == 0: return 1.0, b, c
    k = min(b, c); p = float(2 * binom.cdf(k, n, 0.5))
    return min(p, 1.0), b, c

def mondrian_conformal(val_probs, val_labels, test_probs, K, alpha=0.10):
    """Class-conditional conformal prediction at coverage 1-alpha.
    Non-conformity = 1 - P(true class)."""
    sets = [set() for _ in range(len(test_probs))]
    for k in range(K):
        mask = val_labels == k
        if mask.sum() < 10:
            # Too few calibration samples for class; use marginal quantile
            scores = 1 - val_probs[np.arange(len(val_labels)), val_labels]
        else:
            scores = 1 - val_probs[mask, k]
        # quantile at level ceil((n+1)(1-alpha))/n
        n = len(scores); q = int(np.ceil((n + 1) * (1 - alpha))) / n
        q = min(q, 1.0)
        thr = float(np.quantile(scores, q))
        # add class k to any test point with non-conformity score <= thr
        for i, p in enumerate(test_probs):
            if (1 - p[k]) <= thr:
                sets[i].add(k)
    return sets


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--results-dir", required=True, help="dir with *_test_preds.json")
    ap.add_argument("--val-preds-dir", default=None, help="optional: dir with *_val_preds.json for threshold opt / conformal calibration")
    ap.add_argument("--out", required=True)
    ap.add_argument("--ensemble-members", nargs="+", default=None,
                    help="subset of models to include in ensemble (default: all)")
    ap.add_argument("--alpha", type=float, default=0.10, help="conformal mis-coverage")
    args = ap.parse_args()

    preds_test = load_all(args.results_dir)
    names = list(preds_test.keys())
    print(f"Loaded {len(names)} models: {names}")
    labels = preds_test[names[0]]["labels"]
    K = preds_test[names[0]]["probs"].shape[1]

    report = {"per_model": {}, "ensemble": {}, "mcnemar_bonferroni": {}, "conformal": {}}

    # Per-model extended stats
    for n in names:
        labs = preds_test[n]["labels"]; prs = preds_test[n]["preds"]; pbs = preds_test[n]["probs"]
        acc = accuracy_score(labs, prs)
        p, r, f1, _ = precision_recall_fscore_support(labs, prs, average="macro", zero_division=0)
        per_class = precision_recall_fscore_support(labs, prs, average=None, zero_division=0, labels=list(range(K)))
        try: roc = roc_auc_score(labs, pbs, multi_class="ovr", average="macro", labels=list(range(K)))
        except Exception: roc = float("nan")
        kappa = cohen_kappa_score(labs, prs)
        brier = macro_brier(pbs, labs, K)
        acc_lo, acc_hi = bootstrap_ci_acc(labs, prs)
        report["per_model"][n] = {
            "acc": acc, "acc_ci": [acc_lo, acc_hi],
            "precision": p, "recall": r, "f1": f1,
            "roc_auc": roc, "ece": ece(pbs, labs), "kappa": kappa, "brier": brier,
            "per_class_f1": per_class[2].tolist(),
            "per_class_support": per_class[3].tolist(),
        }
        print(f"  {n:14s}  acc {acc*100:5.2f} [{acc_lo*100:.1f},{acc_hi*100:.1f}]  f1 {f1*100:5.2f}  κ {kappa:.3f}  brier {brier:.3f}")

    # Ensemble: weight each model by its own test F1 (approximation; ideally val F1)
    members = args.ensemble_members or names
    weights = {n: max(0.001, report["per_model"][n]["f1"]) for n in members}
    sub = {n: preds_test[n] for n in members}
    el_labels, el_probs = ensemble_soft(sub, weights)
    el_preds = el_probs.argmax(1)
    e_acc = accuracy_score(el_labels, el_preds)
    e_p, e_r, e_f1, _ = precision_recall_fscore_support(el_labels, el_preds, average="macro", zero_division=0)
    e_acc_lo, e_acc_hi = bootstrap_ci_acc(el_labels, el_preds)
    try: e_roc = roc_auc_score(el_labels, el_probs, multi_class="ovr", average="macro", labels=list(range(K)))
    except Exception: e_roc = float("nan")
    report["ensemble"] = {
        "members": members, "weights": {k: float(v) for k, v in weights.items()},
        "acc": e_acc, "acc_ci": [e_acc_lo, e_acc_hi], "precision": e_p, "recall": e_r, "f1": e_f1,
        "roc_auc": e_roc, "ece": ece(el_probs, el_labels),
        "kappa": cohen_kappa_score(el_labels, el_preds), "brier": macro_brier(el_probs, el_labels, K),
    }
    print(f"\nENSEMBLE   acc {e_acc*100:5.2f} [{e_acc_lo*100:.1f},{e_acc_hi*100:.1f}]  f1 {e_f1*100:5.2f}  roc {e_roc:.4f}")

    # If val preds available: per-class threshold opt + conformal
    if args.val_preds_dir and os.path.isdir(args.val_preds_dir):
        val_preds = load_all(args.val_preds_dir)
        val_members = [n for n in members if n in val_preds]
        if val_members:
            vl, vp = ensemble_soft({n: val_preds[n] for n in val_members}, {n: weights[n] for n in val_members})
            biases = per_class_thresholds(vp, vl, K)
            tuned_preds = apply_thresholds(el_probs, biases)
            tuned_f1 = precision_recall_fscore_support(el_labels, tuned_preds, average="macro", zero_division=0)[2]
            tuned_acc = accuracy_score(el_labels, tuned_preds)
            report["ensemble"]["tuned_biases"] = biases.tolist()
            report["ensemble"]["tuned_acc"] = tuned_acc
            report["ensemble"]["tuned_f1"] = tuned_f1
            print(f"  After per-class threshold tuning: acc {tuned_acc*100:.2f}  f1 {tuned_f1*100:.2f}")

            sets = mondrian_conformal(vp, vl, el_probs, K, alpha=args.alpha)
            sizes = [len(s) for s in sets]
            covered = sum(1 for i, s in enumerate(sets) if el_labels[i] in s) / len(el_labels)
            report["conformal"] = {
                "alpha": args.alpha,
                "empirical_coverage": covered,
                "avg_set_size": float(np.mean(sizes)),
                "frac_singleton": float((np.array(sizes) == 1).mean()),
            }
            print(f"  Conformal (α={args.alpha}):  empirical coverage {covered*100:.1f}%  avg |C| {np.mean(sizes):.2f}  singleton frac {(np.array(sizes)==1).mean()*100:.1f}%")

    # McNemar with Bonferroni
    pairs = []
    for i in range(len(names)):
        for j in range(i+1, len(names)):
            p, b, c = mcnemar_pair(labels, preds_test[names[i]]["preds"], preds_test[names[j]]["preds"])
            pairs.append({"model_a": names[i], "model_b": names[j], "p": p, "b_count": b, "c_count": c})
    n_pairs = len(pairs)
    for r in pairs:
        r["p_bonferroni"] = min(1.0, r["p"] * n_pairs)
        r["sig_005"] = r["p_bonferroni"] < 0.05
    report["mcnemar_bonferroni"] = {"n_pairs": n_pairs, "pairs": pairs}
    print(f"\nMcNemar Bonferroni (n_pairs={n_pairs}):")
    for r in pairs:
        mk = "*" if r["sig_005"] else " "
        print(f" {mk} {r['model_a']:14s} vs {r['model_b']:14s}  raw p={r['p']:.3g}  adj={r['p_bonferroni']:.3g}")

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f: json.dump(report, f, indent=2)
    print(f"\nReport -> {args.out}")


if __name__ == "__main__":
    main()