#!/usr/bin/env python3
"""Ensemble + multi-scale TTA + rejection + optional LM/GNN rescoring.

v2: weighted ensemble (DINOv3 0.45 / ConvNeXt-V2 0.30 / SigLIP2 0.25),
    multi-scale TTA (224/320/384), selective classification @ τ=0.6,
    optional tablet-context GNN post-processing and KenLM rescoring.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter, defaultdict

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from PIL import Image

ROOT = Path("/arf/scratch/stakan/hitit-proje")
sys.path.insert(0, str(ROOT / "hitit_ocr/src"))

from train_classification import build_backbone, get_arch_img_size

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

DEFAULT_WEIGHTS = {
    'dinov3_vitl14': 0.45, 'dinov3_vitb14': 0.45,
    'convnextv2_large': 0.30,
    'siglip2_so400m': 0.25,
}

class ValDataset(Dataset):
    def __init__(self, manifest, label_to_idx, tf, val_fold=0):
        self.records = []
        with open(manifest) as f:
            for line in f:
                r = json.loads(line)
                if r.get('task') != 'classification': continue
                if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
                if not r.get('path') or r.get('storage') != 'fs': continue
                if r.get('integrity_ok') is False: continue
                if (r.get('class_sample_count') or 0) < 5: continue
                if r.get('tablet_view_fold', 0) != val_fold: continue
                self.records.append(r)
        self.label_to_idx = label_to_idx
        self.tf = tf

    def __len__(self): return len(self.records)
    def __getitem__(self, idx):
        r = self.records[idx]
        img = Image.open(r['path']).convert('RGB')
        tablet_id = r.get('tablet_id') or Path(r.get('path', '')).parent.name
        return self.tf(img), self.label_to_idx[r['unified_label']], tablet_id

def collate(batch):
    imgs = torch.stack([b[0] for b in batch])
    ys = torch.tensor([b[1] for b in batch], dtype=torch.long)
    tids = [b[2] for b in batch]
    return imgs, ys, tids

def get_patch_mil_transforms(n_patches=4, native=224,
                              mean=(0.489, 0.448, 0.424), std=(0.362, 0.359, 0.364)):
    """Split image into n_patches × n_patches grid, return list of crops."""
    from torchvision import transforms
    from torchvision.transforms import functional as TF

    class PatchGrid:
        def __init__(self, row, col, n):
            self.row = row; self.col = col; self.n = n
        def __call__(self, img):
            W, H = img.size
            pw, ph = W // self.n, H // self.n
            x0 = self.col * pw; y0 = self.row * ph
            return TF.crop(img, y0, x0, ph, pw)

    # 2x2 patches + full image (5 total)
    out = [('full', transforms.Compose([
        transforms.Resize((native, native), antialias=True),
        transforms.ToTensor(), transforms.Normalize(mean, std)]))]
    for r in range(2):
        for c in range(2):
            out.append((f'p{r}{c}', transforms.Compose([
                PatchGrid(r, c, 2),
                transforms.Resize((native, native), antialias=True),
                transforms.ToTensor(), transforms.Normalize(mean, std)])))
    return out

def get_tta_transforms(scales=(224, 320, 384), native=224, mean=(0.489, 0.448, 0.424),
                       std=(0.362, 0.359, 0.364)):
    """Multi-scale × {identity, gamma+, gamma-, rot+3, rot-3}. No hflip (cuneiform asymmetric).

    native: model's expected input resolution. We resize to `s`, apply aug, then to `native`
    so model always sees its trained resolution."""
    from torchvision import transforms
    from torchvision.transforms import functional as TF

    class Gamma:
        def __init__(self, g): self.g = g
        def __call__(self, img): return TF.adjust_gamma(img, self.g)

    class Rot:
        def __init__(self, a): self.a = a
        def __call__(self, img): return TF.rotate(img, self.a)

    aug_list = [
        ('id', None),
        ('g0.9', Gamma(0.9)),
        ('g1.1', Gamma(1.1)),
        ('rot+3', Rot(3)),
        ('rot-3', Rot(-3)),
    ]
    out = []
    for s in scales:
        for name, pre in aug_list:
            ops = [transforms.Resize((s, s), antialias=True)]
            if pre is not None: ops.append(pre)
            if s != native: ops.append(transforms.Resize((native, native), antialias=True))
            ops += [transforms.ToTensor(), transforms.Normalize(mean, std)]
            out.append((f"{s}_{name}", transforms.Compose(ops)))
    return out

@torch.no_grad()
def infer_softmax(model, loader, device, dtype):
    model.eval()
    all_p, all_y, all_tids = [], [], []
    for x, y, tids in loader:
        x = x.to(device, non_blocking=True)
        with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
            logits = model(x)
        p = F.softmax(logits.float(), dim=-1)
        all_p.append(p.cpu()); all_y.append(y); all_tids.extend(tids)
    return torch.cat(all_p), torch.cat(all_y), all_tids

def fit_temperature(logits, targets, max_iter=100):
    """Fit T on validation — Guo 2017. logits: (N, C). T>0 scales sharpness."""
    logits = logits.double()
    T = torch.ones(1, requires_grad=True, dtype=torch.float64)
    opt = torch.optim.LBFGS([T], lr=0.1, max_iter=max_iter)
    def closure():
        opt.zero_grad()
        loss = F.cross_entropy(logits / T.clamp_min(1e-3), targets)
        loss.backward(); return loss
    opt.step(closure)
    return float(T.item())

def optimize_weights_topk(per_model_probs, targets, k=1, n_iter=80):
    """Nelder-Mead-esque coordinate descent on simplex to maximize top-k accuracy.
    per_model_probs: list of [N, C] tensors (softmax). Returns optimized w."""
    import itertools
    M = len(per_model_probs)
    stacked = torch.stack(per_model_probs).double()  # (M, N, C)
    def topk_acc(w):
        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, dim=-1)
        return sum(targets[i].item() in tk[i].tolist() for i in range(len(targets))) / len(targets)
    # Seed: uniform + heuristic (proportional to own top-1)
    ownacc = np.array([(per_model_probs[i].argmax(-1) == targets).float().mean().item()
                        for i in range(M)])
    seeds = [np.ones(M)/M, ownacc/ownacc.sum()]
    best_w, best_s = seeds[0], topk_acc(seeds[0])
    for s in seeds:
        a = topk_acc(s)
        if a > best_s: best_w, best_s = s.copy(), a
    # Coordinate descent over (i, j) pairs
    step = 0.1
    for it in range(n_iter):
        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_acc(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
    return best_w, best_s

def apply_tablet_gnn(probs, tablet_ids, feat_dim_logit=None, device='cuda'):
    """Refine logits with tablet-level attention. probs (N, C) log-space."""
    try:
        from enhancements.tablet_gnn import TabletContextGAT
    except Exception as e:
        log(f"  tablet_gnn unavailable: {e}"); return probs
    N, C = probs.shape
    # Map string tablet_ids → int
    tid_to_int = {}
    tid_int = torch.zeros(N, dtype=torch.long)
    for i, t in enumerate(tablet_ids):
        if t not in tid_to_int: tid_to_int[t] = len(tid_to_int)
        tid_int[i] = tid_to_int[t]
    # Feed log-probs as "features" (C-dim); randomly init GNN (no pretrained checkpoint).
    # Without trained weights, GNN applies untrained transform — not beneficial.
    # Gate: only apply if a checkpoint exists.
    ckpt = ROOT / 'hitit_ocr/runs/h100/tablet_gnn/ckpt.pt'
    if not ckpt.exists():
        log(f"  tablet_gnn checkpoint absent; skipping (train via enhancements/tablet_gnn.py first)")
        return probs
    gnn = TabletContextGAT(feat_dim=C).to(device)
    gnn.load_state_dict(torch.load(ckpt, map_location=device))
    gnn.eval()
    with torch.no_grad():
        logits_in = torch.log(probs.clamp_min(1e-9)).to(device)
        refined = gnn(logits_in, tid_int.to(device))
        refined_p = F.softmax(refined.float(), dim=-1).cpu()
    return refined_p

def tier_metrics(probs, targets, class_counts):
    """Report per-tier top-1: head (>100), mid (20-100), tail (<20).

    class_counts: dict label_idx → train count.
    """
    pred = probs.argmax(-1)
    tiers = {'head': [], 'mid': [], 'tail': []}
    for i in range(len(targets)):
        c = int(targets[i]); n = class_counts.get(c, 0)
        tier = 'head' if n > 100 else 'mid' if n >= 20 else 'tail'
        tiers[tier].append((int(pred[i]) == c, n))
    out = {}
    for t, rows in tiers.items():
        if not rows: out[t] = {'top1': None, 'n': 0}; continue
        acc = sum(r[0] for r in rows) / len(rows)
        out[t] = {'top1': acc, 'n': len(rows),
                  'median_train_count': sorted([r[1] for r in rows])[len(rows)//2]}
    return out

def multi_view_aggregate(probs, tablet_ids, bbox_hashes=None, mode='mean'):
    """Group probs by (tablet_id, bbox_hash); within each group, average probs."""
    from collections import defaultdict
    groups = defaultdict(list)
    for i, tid in enumerate(tablet_ids):
        key = (tid, bbox_hashes[i] if bbox_hashes else i)
        groups[key].append(i)
    refined = probs.clone()
    for key, idxs in groups.items():
        if len(idxs) <= 1: continue
        if mode == 'mean':
            avg = probs[idxs].mean(0)
        elif mode == 'max':
            avg = probs[idxs].max(0).values
        for i in idxs: refined[i] = avg
    return refined

def selective_metrics(probs, targets, thresholds=(0.5, 0.6, 0.7, 0.8, 0.9)):
    """Return dict of threshold→(selective_acc, coverage)."""
    pred = probs.argmax(-1)
    max_p = probs.max(-1).values
    out = {}
    for t in thresholds:
        keep = max_p >= t
        cov = keep.float().mean().item()
        if keep.sum() == 0:
            out[t] = {'selective_acc': 0.0, 'coverage': 0.0}
            continue
        acc = (pred[keep] == targets[keep]).float().mean().item()
        out[t] = {'selective_acc': acc, 'coverage': cov}
    return out

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpts', nargs='+', required=True,
                    help='arch:path, e.g. dinov3_vitb14:runs/h100/hitit_dinov3_cls_v2/best_ema.pt')
    ap.add_argument('--weights', nargs='+', type=float, default=None,
                    help='Per-ckpt weight; if absent, uses DEFAULT_WEIGHTS by arch')
    ap.add_argument('--manifest', default=str(ROOT / 'datasets/sources/hitit_local/manifest_classification.jsonl'))
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--batch-size', type=int, default=128)
    ap.add_argument('--scales', nargs='+', type=int, default=[224, 320, 384])
    ap.add_argument('--rejection-threshold', type=float, default=0.6)
    ap.add_argument('--use-tablet-gnn', action='store_true')
    ap.add_argument('--optimize-weights', action='store_true',
                    help='Grid/Nelder-Mead weight search on val to max top1')
    ap.add_argument('--temperature-scale', action='store_true',
                    help='Fit per-model temperature on val before ensemble')
    ap.add_argument('--patch-mil', action='store_true',
                    help='Add 2x2 patch-MIL alongside scale TTA')
    ap.add_argument('--multi-view', action='store_true',
                    help='Aggregate probs within same tablet_id group')
    ap.add_argument('--test-fold', type=int, default=None,
                    help='If set, also eval on this held-out fold')
    ap.add_argument('--train-manifest', default=None,
                    help='Train manifest to compute class_counts for tier metrics')
    ap.add_argument('--ensemble-space', choices=['softmax', 'logit', 'logsoftmax'],
                    default='logit',
                    help='Where to average: softmax (old, biased by sharpness), logit (scale-aware)')
    ap.add_argument('--output', required=True)
    ap.add_argument('--dump-probs', default=None,
                    help='If set, save probs/targets/label_to_idx .pt for downstream rescoring')
    args = ap.parse_args()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32

    # Weights
    archs = [s.split(':', 1)[0] for s in args.ckpts]
    if args.weights:
        assert len(args.weights) == len(args.ckpts), "--weights must match --ckpts count"
        w = np.array(args.weights, dtype=np.float64)
    else:
        w = np.array([DEFAULT_WEIGHTS.get(a, 1.0) for a in archs], dtype=np.float64)
    w = w / w.sum()
    log(f"Ensemble weights: {dict(zip(archs, w.tolist()))}")

    first = args.ckpts[0].split(':', 1)[1]
    ck0 = torch.load(first, map_location='cpu', weights_only=False)
    label_to_idx = ck0['label_to_idx']
    n_cls = len(label_to_idx)
    log(f"Classes: {n_cls}")

    per_model_probs = []
    targets = None
    tablet_ids_global = None
    for spec in args.ckpts:
        arch, path = spec.split(':', 1)
        log(f"Loading {arch} from {path}")
        ck = torch.load(path, map_location='cpu', weights_only=False)
        model = build_backbone(arch, n_classes=n_cls).to(device)
        sd = ck['model']
        sd = {k.replace('module.', '', 1) if k.startswith('module.') else k: v for k, v in sd.items()}
        sd = {k.replace('_orig_mod.', '', 1) if k.startswith('_orig_mod.') else k: v for k, v in sd.items()}
        sd = {k: v for k, v in sd.items() if k != 'n_averaged'}
        miss, unexp = model.load_state_dict(sd, strict=False)
        log(f"  missing={len(miss)}, unexpected={len(unexp)}")

        native = get_arch_img_size(arch)
        tta_tfs = get_tta_transforms(scales=tuple(args.scales), native=native)
        if args.patch_mil:
            tta_tfs = tta_tfs + get_patch_mil_transforms(native=native)
        log(f"  native={native}, TTA variants={len(tta_tfs)}")
        tta_probs = []
        for ti, (name, tf) in enumerate(tta_tfs):
            ds = ValDataset(args.manifest, label_to_idx, tf, val_fold=args.val_fold)
            dl = DataLoader(ds, batch_size=args.batch_size, shuffle=False,
                            num_workers=4, pin_memory=True, collate_fn=collate)
            p, y, tids = infer_softmax(model, dl, device, dtype)
            log(f"  TTA {name}: n={len(y)}, top1={(p.argmax(-1)==y).float().mean().item():.4f}")
            tta_probs.append(p)
            if targets is None: targets = y
            if tablet_ids_global is None: tablet_ids_global = tids
        model_avg = torch.stack(tta_probs).mean(0)
        mt1 = (model_avg.argmax(-1) == targets).float().mean().item()
        log(f"{arch} TTA-ensemble top1: {mt1:.4f}")
        per_model_probs.append(model_avg)
        del model; torch.cuda.empty_cache()

    # Weighted ensemble — space matters
    stacked_prob = torch.stack(per_model_probs)  # (M, N, C)
    fitted_T = [1.0] * len(per_model_probs)
    if args.temperature_scale:
        for i, p in enumerate(per_model_probs):
            lg = torch.log(p.clamp_min(1e-12))
            fitted_T[i] = fit_temperature(lg, targets)
            log(f"T[{archs[i]}] = {fitted_T[i]:.3f}")
        # Apply temperature
        per_model_probs = [F.softmax(torch.log(p.clamp_min(1e-12)) / t, dim=-1)
                            for p, t in zip(per_model_probs, fitted_T)]
        stacked_prob = torch.stack(per_model_probs)
    if args.ensemble_space == 'logit':
        # Re-cast softmax→logit via log then re-softmax after weighted avg
        stacked = torch.log(stacked_prob.clamp_min(1e-12))
    elif args.ensemble_space == 'logsoftmax':
        stacked = torch.log(stacked_prob.clamp_min(1e-12)) - \
                  torch.log(stacked_prob.clamp_min(1e-12)).mean(-1, keepdim=True)
    else:
        stacked = stacked_prob

    # Optional weight optimization on val (risk: val-overfit; report both)
    opt_w = None; opt_top1 = None
    if args.optimize_weights:
        opt_w, opt_top1 = optimize_weights_topk(per_model_probs, targets, k=1)
        log(f"Optimized weights: {dict(zip(archs, opt_w.tolist()))} → top1={opt_top1:.4f}")
        w = opt_w  # override for final report

    w_t = torch.tensor(w, dtype=stacked.dtype).view(-1, 1, 1)
    avg = (stacked * w_t).sum(0)
    if args.ensemble_space == 'logit':
        global_probs = F.softmax(avg, dim=-1)
    else:
        global_probs = avg
    top1 = (global_probs.argmax(-1) == targets).float().mean().item()
    _, top5 = global_probs.topk(5, dim=-1)
    top5_acc = sum(targets[i].item() in top5[i].tolist() for i in range(len(targets))) / len(targets)

    log(f"=== WEIGHTED ENSEMBLE ===")
    log(f"Top-1: {top1:.4f}")
    log(f"Top-5: {top5_acc:.4f}")

    # Optional tablet-GNN refinement
    gnn_top1 = None
    if args.use_tablet_gnn and tablet_ids_global is not None:
        log("Applying tablet-context GNN...")
        refined = apply_tablet_gnn(global_probs, tablet_ids_global, device=device)
        gnn_top1 = (refined.argmax(-1) == targets).float().mean().item()
        log(f"After GNN Top-1: {gnn_top1:.4f}")
        global_probs = refined

    # Multi-view aggregation (optional)
    if args.multi_view and tablet_ids_global is not None:
        before = (global_probs.argmax(-1) == targets).float().mean().item()
        global_probs = multi_view_aggregate(global_probs, tablet_ids_global, mode='mean')
        after = (global_probs.argmax(-1) == targets).float().mean().item()
        log(f"Multi-view agg: {before:.4f} → {after:.4f}")

    # Tier metrics (head/mid/tail)
    tier_out = None
    train_manifest = args.train_manifest or args.manifest
    try:
        cls_counts = Counter()
        for line in open(train_manifest):
            r = json.loads(line)
            if r.get('task') != 'classification': continue
            if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
            if r.get('tablet_view_fold', 0) == args.val_fold: continue  # train only
            cls_counts[label_to_idx[r['unified_label']]] += 1
        tier_out = tier_metrics(global_probs, targets, cls_counts)
        log("Tier metrics:")
        for t, m in tier_out.items():
            if m['top1'] is not None:
                log(f"  {t}: top1={m['top1']:.4f} n={m['n']}")
    except Exception as e:
        log(f"tier_metrics failed: {e}")

    # Selective classification
    sel = selective_metrics(global_probs, targets)
    log("Selective metrics:")
    for t, m in sel.items():
        log(f"  τ={t}: sel_acc={m['selective_acc']:.4f}, coverage={m['coverage']:.4f}")

    # Per-model breakdown
    per_model = {}
    for i, spec in enumerate(args.ckpts):
        arch = spec.split(':', 1)[0]
        t1 = (per_model_probs[i].argmax(-1) == targets).float().mean().item()
        per_model[arch] = t1

    # Pairwise agreement (diversity)
    agree = {}
    for i in range(len(per_model_probs)):
        for j in range(i+1, len(per_model_probs)):
            ai = per_model_probs[i].argmax(-1)
            aj = per_model_probs[j].argmax(-1)
            key = f"{archs[i]}_vs_{archs[j]}"
            agree[key] = (ai == aj).float().mean().item()

    out = {
        'ensemble_top1': top1, 'ensemble_top5': top5_acc,
        'ensemble_space': args.ensemble_space,
        'fitted_temperatures': dict(zip(archs, fitted_T)) if args.temperature_scale else None,
        'ensemble_weights': dict(zip(archs, (w.tolist() if isinstance(w, np.ndarray) else list(w)))),
        'optimized_top1': opt_top1,
        'optimized_weights': dict(zip(archs, opt_w.tolist())) if opt_w is not None else None,
        'per_model_tta_top1': per_model,
        'pairwise_agreement': agree,
        'selective_metrics': {str(t): m for t, m in sel.items()},
        'rejection_threshold_default': args.rejection_threshold,
        'gnn_top1': gnn_top1,
        'tier_metrics': tier_out,
        'multi_view': args.multi_view,
        'scales': args.scales,
        'n_val': len(targets), 'n_classes': n_cls,
        'ckpts': args.ckpts,
    }
    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 args.dump_probs:
        Path(args.dump_probs).parent.mkdir(parents=True, exist_ok=True)
        dump = {
            'probs': global_probs.cpu(),
            'targets': targets.cpu(),
            'label_to_idx': label_to_idx,
            'tablet_ids': tablet_ids_global,
            'per_model_probs': [p.cpu() for p in per_model_probs],
            'archs': archs,
        }
        torch.save(dump, args.dump_probs)
        log(f"Dumped probs: {args.dump_probs}")

if __name__ == '__main__':
    main()