#!/usr/bin/env python3
"""Knowledge distillation: ensemble teacher → single DINOv3-L student.

Loss: DIST (Huang et al., NeurIPS 2022) — Pearson correlation across class dim
+ classical KD (KL soft + CE hard).

Usage:
    python distillation.py \
        --teachers dinov3_vitb14:...ema.pt convnextv2_large:... siglip2_so400m:... \
        --teacher-weights 0.45 0.30 0.25 \
        --student-arch dinov3_vitb14 \
        --manifest ... --output runs/h100/distilled_student/
"""
import os, sys, json, argparse, time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torch.optim.swa_utils import AveragedModel
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, HititClsDataset

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

def dist_loss(student_logits, teacher_probs, T=4.0):
    """DIST = KL(teacher || student) + Pearson correlation."""
    # KL (soft)
    s_logp = F.log_softmax(student_logits / T, dim=-1)
    t_p = teacher_probs  # already softmax (ensemble, mixed T)
    t_p_T = F.softmax(torch.log(t_p.clamp_min(1e-12)) / T, dim=-1)
    kd = F.kl_div(s_logp, t_p_T, reduction='batchmean') * (T * T)

    # Inter-class correlation (Pearson over class dim, per-sample)
    def _zscore(x):
        return (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + 1e-6)
    s_z = _zscore(F.softmax(student_logits, dim=-1))
    t_z = _zscore(t_p)
    inter = 1.0 - (s_z * t_z).mean()

    # Intra-sample correlation (Pearson across batch, per-class)
    s_z2 = _zscore(F.softmax(student_logits, dim=-1).T)
    t_z2 = _zscore(t_p.T)
    intra = 1.0 - (s_z2 * t_z2).mean()

    return kd + inter + intra

@torch.no_grad()
def ensemble_teacher_probs(teachers, teacher_ws, x, dtype):
    probs = []
    for m in teachers:
        with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
            logits = m(x)
        probs.append(F.softmax(logits.float(), dim=-1))
    stacked = torch.stack(probs)
    w = torch.tensor(teacher_ws, dtype=stacked.dtype, device=stacked.device).view(-1, 1, 1)
    return (stacked * w).sum(0)

def load_model(arch, ckpt_path, n_cls, device):
    model = build_backbone(arch, n_classes=n_cls).to(device)
    ck = torch.load(ckpt_path, map_location='cpu', weights_only=False)
    sd = ck['model']
    sd = {k.replace('module.', '', 1): v for k, v in sd.items()}
    sd = {k.replace('_orig_mod.', '', 1): v for k, v in sd.items()}
    sd = {k: v for k, v in sd.items() if k != 'n_averaged'}
    model.load_state_dict(sd, strict=False)
    return model, ck

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--teachers', nargs='+', required=True, help='arch:path ...')
    ap.add_argument('--teacher-weights', nargs='+', type=float, default=None)
    ap.add_argument('--student-arch', default='dinov3_vitb14')
    ap.add_argument('--manifest', default=str(ROOT / 'datasets/sources/hitit_local/manifest_classification.jsonl'))
    ap.add_argument('--output', required=True)
    ap.add_argument('--epochs', type=int, default=30)
    ap.add_argument('--batch-size', type=int, default=64)
    ap.add_argument('--lr', type=float, default=1e-4)
    ap.add_argument('--T', type=float, default=4.0)
    ap.add_argument('--hard-weight', type=float, default=0.3)
    ap.add_argument('--val-fold', type=int, default=0)
    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

    first_arch, first_path = args.teachers[0].split(':', 1)
    ck0 = torch.load(first_path, map_location='cpu', weights_only=False)
    label_to_idx = ck0['label_to_idx']
    n_cls = len(label_to_idx)

    # Load teachers (frozen)
    teachers = []
    for t in args.teachers:
        a, p = t.split(':', 1)
        m, _ = load_model(a, p, n_cls, device)
        for param in m.parameters(): param.requires_grad = False
        m.eval()
        teachers.append(m)
    tw = args.teacher_weights or [1.0/len(teachers)]*len(teachers)
    s = sum(tw); tw = [x/s for x in tw]
    log(f"Teachers: {[t.split(':', 1)[0] for t in args.teachers]} weights={tw}")

    # Student — initialize from strongest teacher of same arch if exists, else fresh
    student, _ = None, None
    for t in args.teachers:
        a, p = t.split(':', 1)
        if a == args.student_arch:
            student, _ = load_model(a, p, n_cls, device)
            log(f"Student init from teacher {a}:{p}"); break
    if student is None:
        student = build_backbone(args.student_arch, n_classes=n_cls).to(device)
        log(f"Student fresh: {args.student_arch}")

    # Data — use train split (not val)
    import yaml
    cfg = {'training': {'batch_size': args.batch_size}}
    train_ds = HititClsDataset(args.manifest, cfg, is_train=True,
                                val_fold=args.val_fold, label_to_idx=label_to_idx)
    val_ds = HititClsDataset(args.manifest, cfg, is_train=False,
                              val_fold=args.val_fold, label_to_idx=label_to_idx)
    tr_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
                        num_workers=6, pin_memory=True, drop_last=True)
    va_dl = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=4, pin_memory=True)

    opt = torch.optim.AdamW(student.parameters(), lr=args.lr, weight_decay=0.05)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

    ema = AveragedModel(student, avg_fn=lambda a, n, _: 0.9995*a + 0.0005*n)

    out_dir = Path(args.output); out_dir.mkdir(parents=True, exist_ok=True)
    best_acc = 0.0
    for ep in range(args.epochs):
        student.train()
        tloss, nb = 0.0, 0
        for x, y in tr_dl:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            t_probs = ensemble_teacher_probs(teachers, tw, x, dtype)
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                s_logits = student(x)
                loss_kd = dist_loss(s_logits, t_probs, T=args.T)
                loss_ce = F.cross_entropy(s_logits, y, label_smoothing=0.1)
                loss = (1 - args.hard_weight) * loss_kd + args.hard_weight * loss_ce
            opt.zero_grad(set_to_none=True); loss.backward()
            torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
            opt.step()
            ema.update_parameters(student)
            tloss += loss.item(); nb += 1
        sched.step()

        # Validate EMA
        ema.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for x, y in va_dl:
                x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
                with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                    logits = ema(x)
                correct += (logits.argmax(-1) == y).sum().item(); total += y.size(0)
        acc = correct / max(1, total)
        log(f"ep {ep+1}/{args.epochs}: loss={tloss/max(1,nb):.4f} val_acc={acc:.4f}")
        if acc > best_acc:
            best_acc = acc
            torch.save({'model': ema.state_dict(), 'label_to_idx': label_to_idx,
                        'arch': args.student_arch, 'acc': acc},
                       out_dir / 'best_ema.pt')
    log(f"BEST student acc: {best_acc:.4f} → {out_dir / 'best_ema.pt'}")

if __name__ == '__main__':
    main()