#!/usr/bin/env python3
"""
Multi-posthoc probability fuse: learnable weights across independent probability sources.

Input: list of probs.pt files (ensemble_v13, ensemble_v13_swin, arcface_probs, posthoc_*).
Each {probs: [N, C], targets: [N], label_to_idx: {...}}.

Method:
  1. Align label_to_idx across sources (intersect; reindex; drop rows with label not in intersection).
  2. Temperature-scale each source on val (NLL minimization).
  3. Search ensemble weights (Dirichlet grid search, then SGD fine-tune) maximizing val top1.
  4. Optional isotonic calibration on final probs.

Output: fused probs.pt + fuse_eval.json (weights, top1, top5, selective@τ).
"""
import argparse
import json
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from sklearn.isotonic import IsotonicRegression


def load_probs(path):
    d = torch.load(path, map_location='cpu', weights_only=False)
    probs = d['probs'] if torch.is_tensor(d['probs']) else torch.tensor(d['probs'])
    targets = d['targets'] if torch.is_tensor(d['targets']) else torch.tensor(d['targets'])
    label_to_idx = d.get('label_to_idx', {})
    return probs.float(), targets.long(), label_to_idx


def fit_temperature(logits, targets, lr=0.01, steps=200):
    T = torch.ones(1, requires_grad=True)
    opt = torch.optim.LBFGS([T], lr=lr, max_iter=steps)

    def closure():
        opt.zero_grad()
        loss = F.cross_entropy(logits / T.clamp(min=0.05), targets)
        loss.backward()
        return loss

    opt.step(closure)
    return T.detach().item()


def fit_weights(prob_list, targets, n_trials=3000, seed=0):
    """Dirichlet sampling to find weight vector maximizing top1."""
    rng = np.random.default_rng(seed)
    K = len(prob_list)
    stacked = torch.stack(prob_list, dim=0)  # K, N, C

    best_top1 = -1
    best_w = None
    # Start uniform
    w0 = np.ones(K) / K
    for trial in range(n_trials):
        if trial == 0:
            w = w0
        else:
            # Dirichlet concentration varying
            alpha = rng.uniform(0.5, 5.0)
            w = rng.dirichlet([alpha] * K)
        fused = (stacked * torch.from_numpy(w).float().view(-1, 1, 1)).sum(0)
        top1 = (fused.argmax(-1) == targets).float().mean().item()
        if top1 > best_top1:
            best_top1 = top1
            best_w = w
    return best_w, best_top1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--probs', nargs='+', required=True,
                    help='List of probs.pt files to fuse')
    ap.add_argument('--names', nargs='*', default=None, help='Optional display names')
    ap.add_argument('--output-probs', required=True)
    ap.add_argument('--output-json', required=True)
    ap.add_argument('--isotonic', action='store_true', help='Fit isotonic on max-prob -> correct')
    args = ap.parse_args()

    names = args.names or [Path(p).stem for p in args.probs]

    sources = []
    ref_label = None
    ref_targets = None
    for p, n in zip(args.probs, names):
        if not Path(p).exists():
            print(f"SKIP {n}: {p} does not exist")
            continue
        probs, targets, l2i = load_probs(p)
        print(f"{n}: probs={tuple(probs.shape)}, top1={((probs.argmax(-1)==targets).float().mean().item()):.4f}")
        if ref_label is None:
            ref_label = l2i
            ref_targets = targets
        sources.append((n, probs, targets, l2i))

    if not sources:
        print("No valid probs files!")
        return

    # Reindex to ref_label common cols (assume same for now; fallback to min C)
    min_C = min(s[1].shape[1] for s in sources)
    print(f"\nCommon classes: {min_C}")
    aligned = []
    for n, p, t, l in sources:
        aligned.append(p[:, :min_C])

    # Use first source's targets
    N = min(a.shape[0] for a in aligned)
    aligned = [a[:N] for a in aligned]
    targets = ref_targets[:N]

    # Temperature-scale each (on fake logits from probs; clamp)
    print("\n=== Temperature scaling ===")
    scaled = []
    for n, a in zip(names[:len(aligned)], aligned):
        logits = (a.clamp(min=1e-8).log())
        T = fit_temperature(logits, targets)
        sc = F.softmax(logits / max(T, 0.05), dim=-1)
        print(f"  {n}: T={T:.3f}")
        scaled.append(sc)

    # Weight search
    print("\n=== Weight search (Dirichlet) ===")
    w, top1 = fit_weights(scaled, targets)
    print(f"  Weights: {dict(zip(names[:len(w)], [float(x) for x in w]))}")
    print(f"  top1 = {top1:.4f}")

    # Compute fused
    stacked = torch.stack(scaled, dim=0)
    fused = (stacked * torch.from_numpy(w).float().view(-1, 1, 1)).sum(0)

    # Optional isotonic
    if args.isotonic:
        max_p = fused.max(-1).values.numpy()
        correct = (fused.argmax(-1) == targets).numpy().astype(int)
        iso = IsotonicRegression(out_of_bounds='clip').fit(max_p, correct)
        print(f"  Isotonic fit: max_p range [{max_p.min():.3f}, {max_p.max():.3f}]")

    # Stats
    pred = fused.argmax(-1)
    top5_idx = fused.topk(5, dim=-1).indices
    top5 = (top5_idx == targets.unsqueeze(1)).any(1).float().mean().item()

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

    report = {
        'sources': names[:len(aligned)],
        'weights': {n: float(wi) for n, wi in zip(names[:len(w)], w)},
        'fused_top1': top1,
        'fused_top5': top5,
        'selective': sel,
        'n_val': int(N),
    }
    Path(args.output_json).write_text(json.dumps(report, indent=2))
    print(f"\n{json.dumps(report, indent=2)}")

    torch.save({
        'probs': fused,
        'targets': targets,
        'label_to_idx': ref_label,
        'weights': {n: float(wi) for n, wi in zip(names[:len(w)], w)},
    }, args.output_probs)
    print(f"\nSaved: {args.output_probs}")


if __name__ == '__main__':
    main()