#!/usr/bin/env python3
"""Fast distillation with pre-computed teacher logits cache.

Two phases:
  Phase 1: Teacher inference pass — compute weighted ensemble probs for every
  train sample (no augmentation, single forward). Save to disk.

  Phase 2: Fast student training — load cached teacher probs, pair with
  augmented student input, KD loss. No teacher forward in loop → 10-20× faster.

Augmentation note: teacher probs are computed on the deterministic
transformation (resize+normalize, no mixup). Student still uses mixup etc.
This is standard DIST/KD practice (teacher on clean view).
"""
import os, sys, json, argparse, time
from pathlib import Path

import numpy as np
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, get_arch_img_size

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

class TeacherCacheBuilder:
    """Build a per-sample ensemble teacher probs tensor."""
    def __init__(self, teachers, weights, device, dtype):
        self.teachers = teachers; self.ws = weights
        self.device = device; self.dtype = dtype

    @torch.no_grad()
    def build(self, ds, batch_size=128):
        if len(ds) == 0:
            log(f"  ERROR: dataset empty"); return None
        loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                            num_workers=4, pin_memory=True)
        log(f"  Building teacher cache on {len(ds)} samples...")
        all_probs = None
        for i, batch in enumerate(loader):
            x = batch[0].to(self.device, non_blocking=True)
            probs = None
            for m, w in zip(self.teachers, self.ws):
                m.eval()
                with torch.amp.autocast('cuda', dtype=self.dtype, enabled=True):
                    lg = m(x)
                p = F.softmax(lg.float(), dim=-1)
                probs = p * w if probs is None else probs + p * w
            probs = probs / sum(self.ws)
            all_probs = probs.cpu() if all_probs is None else torch.cat([all_probs, probs.cpu()])
            if (i+1) % 20 == 0:
                log(f"  teacher cache: {all_probs.size(0)}/{len(ds)}")
        log(f"  Cache built: {all_probs.shape if all_probs is not None else 'None'}")
        return all_probs

class CachedDistillDataset(Dataset):
    """Train dataset returning (augmented_image, hard_label, cached_teacher_probs)."""
    def __init__(self, inner_ds, teacher_probs):
        self.ds = inner_ds
        self.probs = teacher_probs
        assert len(self.ds) == len(self.probs), (len(self.ds), len(self.probs))
    def __len__(self): return len(self.ds)
    def __getitem__(self, i):
        img, y = self.ds[i]
        return img, y, self.probs[i]

def dist_loss(student_logits, teacher_probs, T=4.0):
    s_logp = F.log_softmax(student_logits / T, dim=-1)
    t_p_T = F.softmax(torch.log(teacher_probs.clamp_min(1e-12)) / T, dim=-1)
    kd = F.kl_div(s_logp, t_p_T, reduction='batchmean') * (T * T)
    # Pearson intra-sample
    def _zscore(x): return (x - x.mean(-1, keepdim=True)) / (x.std(-1, keepdim=True) + 1e-6)
    s_z = _zscore(F.softmax(student_logits, dim=-1)); t_z = _zscore(teacher_probs)
    inter = 1.0 - (s_z * t_z).mean()
    return kd + inter

def load_teacher(arch, path, n_cls, device):
    model = build_backbone(arch, n_classes=n_cls).to(device)
    ck = torch.load(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)
    for p in model.parameters(): p.requires_grad = False
    model.eval()
    return model, ck

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--teachers', nargs='+', required=True)
    ap.add_argument('--teacher-weights', nargs='+', type=float, default=None)
    ap.add_argument('--student-arch', default='dinov3_vitl14')
    ap.add_argument('--manifest', default=str(ROOT / 'datasets/sources/hitit_local/manifest_classification_stratified_aug.jsonl'))
    ap.add_argument('--output', required=True)
    ap.add_argument('--epochs', type=int, default=20)
    ap.add_argument('--batch-size', type=int, default=96)
    ap.add_argument('--lr', type=float, default=5e-5)
    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)
    ap.add_argument('--min-samples', type=int, default=10)
    ap.add_argument('--cache-only', action='store_true',
                    help='Build cache then exit (useful for separate jobs)')
    ap.add_argument('--cache-path', default=None,
                    help='Where to store/load teacher cache (default: output/teacher_cache.pt)')
    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)
    tw = args.teacher_weights or [1.0]*len(args.teachers)
    s = sum(tw); tw = [x/s for x in tw]

    out_dir = Path(args.output); out_dir.mkdir(parents=True, exist_ok=True)
    cache_path = Path(args.cache_path) if args.cache_path else (out_dir / 'teacher_cache.pt')

    # Dataset shared by teacher (deterministic tf) and student (augmented)
    import yaml
    cfg = {'training': {'batch_size': args.batch_size}, 'img_size': get_arch_img_size(args.student_arch)}

    # Build TRAIN split with deterministic transform for teacher caching
    teacher_ds = HititClsDataset(args.manifest,
                                  {'training': {'batch_size': args.batch_size},
                                   'img_size': get_arch_img_size(first_arch)},
                                  is_train=True,
                                  val_fold=args.val_fold,
                                  label_to_idx=label_to_idx,
                                  min_samples=args.min_samples)
    # Replace train augmentation with deterministic eval transform
    from torchvision import transforms as _tv
    _img_size = get_arch_img_size(first_arch)
    teacher_ds.tf = _tv.Compose([
        _tv.Resize((_img_size, _img_size), antialias=True),
        _tv.ToTensor(),
        _tv.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
    ])
    log(f"Teacher cache: {len(teacher_ds)} train records")

    # Build / load teacher cache (validate not None)
    teacher_probs = None
    if cache_path.exists():
        log(f"Loading teacher cache from {cache_path}")
        try:
            teacher_probs = torch.load(cache_path, map_location='cpu', weights_only=False)
            if teacher_probs is None or not hasattr(teacher_probs, 'size'):
                log("Cache corrupt (None or invalid); rebuilding")
                cache_path.unlink(missing_ok=True); teacher_probs = None
        except Exception as e:
            log(f"Cache load failed: {e}; rebuilding")
            cache_path.unlink(missing_ok=True); teacher_probs = None
    if teacher_probs is None:
        log("Loading teachers...")
        teachers = []
        for t in args.teachers:
            a, p = t.split(':', 1)
            m, _ = load_teacher(a, p, n_cls, device)
            teachers.append(m)
        builder = TeacherCacheBuilder(teachers, tw, device, dtype)
        teacher_probs = builder.build(teacher_ds, batch_size=args.batch_size)
        if teacher_probs is None:
            log("ERROR: teacher cache build returned None; aborting")
            sys.exit(1)
        torch.save(teacher_probs, cache_path)
        log(f"Cache saved: {cache_path} shape={teacher_probs.shape}")
        del teachers; torch.cuda.empty_cache()

    if args.cache_only: return

    # Student training — uses augmented train dataset paired with cached probs
    student_cfg = dict(cfg)
    student_cfg['img_size'] = get_arch_img_size(args.student_arch)
    train_ds = HititClsDataset(args.manifest, student_cfg, is_train=True,
                                val_fold=args.val_fold, label_to_idx=label_to_idx,
                                min_samples=args.min_samples)
    assert len(train_ds) == teacher_probs.size(0), \
        f"train_ds={len(train_ds)} vs cache={teacher_probs.size(0)}"
    cached_ds = CachedDistillDataset(train_ds, teacher_probs)

    val_ds = HititClsDataset(args.manifest, student_cfg, is_train=False,
                              val_fold=args.val_fold, label_to_idx=label_to_idx,
                              min_samples=args.min_samples)

    tr_dl = DataLoader(cached_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)

    # Student init — if teacher has same arch, start from its weights
    student = None
    for t in args.teachers:
        a, p = t.split(':', 1)
        if a == args.student_arch:
            student, _ = load_teacher(a, p, n_cls, device)
            for pp in student.parameters(): pp.requires_grad = True
            log(f"Student init from {a}:{p}")
            break
    if student is None:
        student = build_backbone(args.student_arch, n_classes=n_cls).to(device)

    opt = torch.optim.AdamW(student.parameters(), lr=args.lr, weight_decay=0.05, fused=False)
    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)

    best = 0.0
    for ep in range(args.epochs):
        student.train()
        tl, nb = 0.0, 0
        for x, y, tp in tr_dl:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            tp = tp.to(device, non_blocking=True)
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                s_logits = student(x)
                loss_kd = dist_loss(s_logits, tp, 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(); loss.backward()
            torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
            opt.step(); ema.update_parameters(student)
            tl += loss.item(); nb += 1
        sched.step()
        ema.eval()
        with torch.no_grad():
            correct, total = 0, 0
            for x, y in va_dl:
                x = x.to(device); y = y.to(device)
                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={tl/max(1,nb):.4f} val_acc={acc:.4f}")
        if acc > best:
            best = 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: {best:.4f} → {out_dir / 'best_ema.pt'}")

if __name__ == '__main__':
    main()