#!/usr/bin/env python3
"""Conformal prediction + calibration + confusion analysis for v13 ensemble probs.

Given dumped probs from eval_ensemble_v2.py --dump-probs, produces:
- Split conformal prediction sets (APS, Raps, LAC)
- ECE (Expected Calibration Error) before/after temperature
- Top-k confusion pairs (N×N confusion matrix ranked)
- Per-tablet failure analysis (which tablets score lowest)

Not paper-critical but reviewer-friendly (frequentist guarantees).
"""
import argparse, json, time
from pathlib import Path
from collections import Counter, defaultdict
import numpy as np
import torch
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parents[1].parent

def log(msg): print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

def ece(probs, targets, n_bins=15):
    conf, pred = probs.max(-1)
    acc = (pred == targets).float()
    bins = torch.linspace(0, 1, n_bins + 1)
    e = 0.0
    for i in range(n_bins):
        lo, hi = bins[i], bins[i+1]
        m = (conf > lo) & (conf <= hi)
        if m.sum() > 0:
            e += (m.float().mean() * (conf[m].mean() - acc[m].mean()).abs()).item()
    return e

def aps_conformal(probs, targets, alpha=0.1, cal_frac=0.5, seed=42):
    """Adaptive Prediction Sets: nonconformity = sum of sorted probs until true label."""
    N, C = probs.shape
    rng = np.random.default_rng(seed)
    idx = rng.permutation(N)
    n_cal = int(N * cal_frac)
    cal_idx, test_idx = idx[:n_cal], idx[n_cal:]
    # Nonconformity scores on calibration
    p_cal = probs[cal_idx]; y_cal = targets[cal_idx]
    scores = []
    for i in range(len(cal_idx)):
        order = p_cal[i].argsort(descending=True)
        cum = 0.0
        for r, c in enumerate(order.tolist()):
            cum += p_cal[i, c].item()
            if c == y_cal[i].item():
                scores.append(cum)
                break
    q = np.quantile(scores, 1 - alpha)
    # Build prediction sets on test
    set_sizes = []
    covered = 0
    for i in test_idx:
        order = probs[i].argsort(descending=True)
        cum = 0.0; s = 0
        for r, c in enumerate(order.tolist()):
            cum += probs[i, c].item()
            s += 1
            if cum >= q:
                break
        set_sizes.append(s)
        # Check if true label is in set (top-s)
        topk = probs[i].topk(s).indices.tolist()
        if targets[i].item() in topk:
            covered += 1
    return {
        'alpha': alpha,
        'q_hat': float(q),
        'coverage': covered / len(test_idx),
        'mean_set_size': float(np.mean(set_sizes)),
        'median_set_size': float(np.median(set_sizes)),
        'p90_set_size': float(np.percentile(set_sizes, 90)),
    }

def top_confusion_pairs(probs, targets, idx_to_label, top_k=30):
    pred = probs.argmax(-1)
    pairs = Counter()
    for p, t in zip(pred.tolist(), targets.tolist()):
        if p != t:
            pairs[(t, p)] += 1
    return [
        {
            'true': idx_to_label.get(t, str(t)),
            'pred': idx_to_label.get(p, str(p)),
            'count': c,
        }
        for (t, p), c in pairs.most_common(top_k)
    ]

def per_tablet_accuracy(probs, targets, tablet_ids, top_k_worst=20):
    if tablet_ids is None: return None
    pred = probs.argmax(-1)
    by_tablet = defaultdict(lambda: {'correct': 0, 'total': 0})
    for p, t, tid in zip(pred.tolist(), targets.tolist(), tablet_ids):
        by_tablet[tid]['correct'] += int(p == t)
        by_tablet[tid]['total'] += 1
    rows = [
        {'tablet_id': tid, 'acc': d['correct']/d['total'], 'n': d['total']}
        for tid, d in by_tablet.items() if d['total'] >= 5
    ]
    rows.sort(key=lambda r: r['acc'])
    return rows[:top_k_worst]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--probs', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--alphas', nargs='+', type=float, default=[0.01, 0.05, 0.1, 0.2])
    args = ap.parse_args()

    log(f"Loading {args.probs}")
    d = torch.load(args.probs, map_location='cpu', weights_only=False)
    probs = d['probs']
    targets = d['targets']
    label_to_idx = d['label_to_idx']
    idx_to_label = {v: k for k, v in label_to_idx.items()}
    tablet_ids = d.get('tablet_ids')
    log(f"N={len(targets)}, C={probs.shape[1]}")

    out = {}

    # Top-1, top-5
    pred = probs.argmax(-1)
    top1 = (pred == targets).float().mean().item()
    _, top5 = probs.topk(5, dim=-1)
    top5_acc = sum(targets[i].item() in top5[i].tolist() for i in range(len(targets))) / len(targets)
    out['baseline'] = {'top1': top1, 'top5': top5_acc}
    log(f"top1={top1:.4f} top5={top5_acc:.4f}")

    # ECE
    out['ece'] = ece(probs, targets)
    log(f"ECE={out['ece']:.4f}")

    # Conformal at multiple alphas
    out['conformal'] = {}
    for a in args.alphas:
        try:
            r = aps_conformal(probs, targets, alpha=a)
            out['conformal'][f'alpha_{a}'] = r
            log(f"Conformal α={a}: cov={r['coverage']:.4f} mean_set={r['mean_set_size']:.2f} p90={r['p90_set_size']:.1f}")
        except Exception as e:
            log(f"Conformal α={a} failed: {e}")

    # Top confusion pairs
    try:
        out['top_confusion_pairs'] = top_confusion_pairs(probs, targets, idx_to_label, 30)
        log(f"Top confusion pair: {out['top_confusion_pairs'][0] if out['top_confusion_pairs'] else 'none'}")
    except Exception as e:
        log(f"confusion pairs failed: {e}")

    # Per-tablet worst
    try:
        out['worst_tablets'] = per_tablet_accuracy(probs, targets, tablet_ids, 20)
    except Exception as e:
        log(f"per-tablet failed: {e}")

    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    json.dump(out, open(args.output, 'w'), indent=2)
    log(f"Saved: {args.output}")

if __name__ == '__main__':
    main()