#!/usr/bin/env python3
"""Mega ensemble — union all trained heads by probabilistic averaging.

Inputs: list of pt files, each containing {probs, targets, top1}.
Optimizes weights via coordinate descent on val top1.
"""
import json, argparse, itertools
from pathlib import Path
import torch
import torch.nn.functional as F
import numpy as np


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--probs', nargs='+', required=True, help='name:path pt files')
    ap.add_argument('--optimize', action='store_true')
    ap.add_argument('--output', required=True)
    args = ap.parse_args()

    names, probs, targets = [], [], None
    own_acc = []
    for spec in args.probs:
        if ':' in spec:
            nm, pth = spec.split(':', 1)
        else:
            nm = Path(spec).stem; pth = spec
        if not Path(pth).exists():
            print(f"skip missing {pth}"); continue
        d = torch.load(pth, map_location='cpu', weights_only=False)
        p = d.get('probs')
        if p is None: p = d.get('refined_probs')
        t = d.get('targets')
        if p is None or t is None:
            print(f"skip {pth}: no probs/targets"); continue
        if targets is None:
            targets = t
        elif not torch.equal(t, targets):
            print(f"{nm}: target mismatch (size {t.size(0)} vs {targets.size(0)}); skip")
            continue
        probs.append(p)
        names.append(nm)
        acc = (p.argmax(-1) == targets).float().mean().item()
        own_acc.append(acc)
        print(f"{nm}: top1={acc:.4f}")

    if not probs:
        print("No probs to ensemble"); return

    stacked = torch.stack(probs).double()  # (M, N, C)
    M = stacked.size(0)

    def topk(w, k=1):
        w = torch.tensor(w, dtype=stacked.dtype).view(-1, 1, 1)
        p = (stacked * w).sum(0)
        if k == 1: return (p.argmax(-1) == targets).float().mean().item()
        _, tk = p.topk(k, -1)
        return sum(targets[i].item() in tk[i].tolist() for i in range(len(targets))) / len(targets)

    # Seed: own-acc proportional
    own = np.array(own_acc, dtype=np.float64)
    w_uni = np.ones(M) / M
    w_own = own / own.sum()
    best_w = w_own if topk(w_own) > topk(w_uni) else w_uni
    best_s = max(topk(w_own), topk(w_uni))
    print(f"Seed top1: uniform={topk(w_uni):.4f}, own-acc={topk(w_own):.4f}")

    if args.optimize:
        step = 0.1
        for it in range(100):
            improved = False
            for i, j in itertools.combinations(range(M), 2):
                for d in (step, -step):
                    w = best_w.copy()
                    w[i] += d; w[j] -= d
                    if w.min() < 0: continue
                    w = w / w.sum()
                    a = topk(w)
                    if a > best_s + 1e-6:
                        best_w, best_s = w, a; improved = True
            if not improved: step *= 0.5
            if step < 1e-4: break
        print(f"Optimized top1: {best_s:.4f}")

    # Final
    w_t = torch.tensor(best_w, dtype=stacked.dtype).view(-1, 1, 1)
    final_probs = (stacked * w_t).sum(0)
    top1 = (final_probs.argmax(-1) == targets).float().mean().item()
    _, top5 = final_probs.topk(5, -1)
    top5_acc = sum(targets[i].item() in top5[i].tolist() for i in range(len(targets))) / len(targets)

    # Selective metrics
    max_p = final_probs.max(-1).values
    pred = final_probs.argmax(-1)
    sel = {}
    for thr in (0.5, 0.6, 0.7, 0.8, 0.9):
        keep = max_p >= thr
        cov = keep.float().mean().item()
        acc = (pred[keep] == targets[keep]).float().mean().item() if keep.sum() else 0.0
        sel[str(thr)] = {'selective_acc': acc, 'coverage': cov}

    print(f"=== MEGA ENSEMBLE ===")
    print(f"Top-1: {top1:.4f}, Top-5: {top5_acc:.4f}")
    for t, m in sel.items():
        print(f"  τ={t}: sel_acc={m['selective_acc']:.4f} cov={m['coverage']:.4f}")

    out = {
        'weights': dict(zip(names, best_w.tolist())),
        'top1': top1, 'top5': top5_acc,
        'selective_metrics': sel,
        'per_model_top1': dict(zip(names, own_acc)),
        'n_val': len(targets),
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    json.dump(out, open(args.output, 'w'), indent=2)
    torch.save({'probs': final_probs, 'targets': targets},
                args.output.replace('.json', '_probs.pt'))
    print(f"Saved → {args.output}")

if __name__ == '__main__':
    main()