#!/usr/bin/env python3
"""Advanced training heads with new objectives.

Each head trained on top of FROZEN v4 DINOv3-L features:
  - dual-branch: global + center-zoom
  - 180° rotation symmetry consistency
  - multi-scale (3-view) contrastive
  - phoneme multi-label auxiliary
  - OHEM (loss-weighted batches)

Each produces a probs pt; feed to mega_ensemble.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
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(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)


# ---------------- Dual-branch data ----------------
class DualBranchDS(Dataset):
    """Return (global_224, center_zoom_112_to_224, label)."""
    def __init__(self, manifest, label_to_idx, img_size, is_train, val_fold, min_samples):
        from torchvision import transforms
        cc = Counter()
        self.records = []
        for line in open(manifest):
            r = json.loads(line)
            if r.get('task') != 'classification' or not r.get('unified_label'): continue
            cc[r['unified_label']] += 1
        for line in open(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 not r.get('path') or r.get('storage') != 'fs': continue
            if r.get('integrity_ok') is False: continue
            if cc[r['unified_label']] < min_samples: continue
            fold = r.get('tablet_view_fold', 0)
            if is_train and fold == val_fold: continue
            if not is_train and fold != val_fold: continue
            self.records.append(r)
        self.l2i = label_to_idx
        self.mean = [0.489, 0.448, 0.424]; self.std = [0.362, 0.359, 0.364]
        self.tf_global = transforms.Compose([
            transforms.Resize((img_size, img_size), antialias=True),
            transforms.ToTensor(), transforms.Normalize(self.mean, self.std),
        ])
        # Center crop (inner 50%) then resize back
        self.tf_zoom = transforms.Compose([
            transforms.Resize((img_size, img_size), antialias=True),
            transforms.CenterCrop(img_size // 2),
            transforms.Resize((img_size, img_size), antialias=True),
            transforms.ToTensor(), transforms.Normalize(self.mean, self.std),
        ])
    def __len__(self): return len(self.records)
    def __getitem__(self, i):
        r = self.records[i]
        img = Image.open(r['path']).convert('RGB')
        return self.tf_global(img), self.tf_zoom(img), self.l2i[r['unified_label']]


class DualBranchHead(nn.Module):
    def __init__(self, feat_dim, n_classes):
        super().__init__()
        self.fuse = nn.Sequential(
            nn.Linear(feat_dim * 2, feat_dim), nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(feat_dim, n_classes))
    def forward(self, f_g, f_z):
        return self.fuse(torch.cat([f_g, f_z], dim=-1))


@torch.no_grad()
def extract(bb, x, dtype):
    raw = bb.module if hasattr(bb, 'module') else bb
    with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
        if hasattr(raw, 'forward_features'):
            f = raw.forward_features(x)
            if f.dim() == 3: f = f[:, 0]
            elif f.dim() == 4: f = f.mean(dim=(2, 3))
        else: f = raw(x)
    return f.float()


def train_dual_branch(bb, label_to_idx, manifest, args, device, dtype):
    img_size = get_arch_img_size(args.backbone_arch)
    tr_ds = DualBranchDS(manifest, label_to_idx, img_size, True, args.val_fold, args.min_samples)
    va_ds = DualBranchDS(manifest, label_to_idx, img_size, False, args.val_fold, args.min_samples)
    log(f"DualBranch Train={len(tr_ds)} Val={len(va_ds)}")
    cc = Counter(tr_ds.l2i[r['unified_label']] for r in tr_ds.records)
    sw = np.array([1.0 / max(1, cc[tr_ds.l2i[r['unified_label']]])**0.5
                    for r in tr_ds.records], dtype=np.float32)
    tr_dl = DataLoader(tr_ds, batch_size=args.batch_size,
                        sampler=WeightedRandomSampler(sw, len(tr_ds), True),
                        num_workers=6, pin_memory=True, drop_last=True)
    va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=4, pin_memory=True)

    with torch.no_grad():
        feat_dim = extract(bb, torch.zeros(1, 3, img_size, img_size, device=device), dtype).size(-1)
    head = DualBranchHead(feat_dim, len(label_to_idx)).to(device)
    opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

    best, best_probs = 0, None
    for ep in range(args.epochs):
        head.train()
        tl, nb = 0, 0
        for xg, xz, y in tr_dl:
            xg = xg.to(device); xz = xz.to(device); y = y.to(device)
            fg = extract(bb, xg, dtype); fz = extract(bb, xz, dtype)
            logits = head(fg, fz)
            # OHEM: weight by loss
            per_loss = F.cross_entropy(logits, y, reduction='none', label_smoothing=0.1)
            w_ohem = per_loss.detach()
            w_ohem = w_ohem / w_ohem.mean().clamp_min(1e-6)
            loss = (per_loss * w_ohem).mean()
            opt.zero_grad(); loss.backward(); opt.step()
            tl += loss.item(); nb += 1
        sched.step()
        head.eval()
        cs, tot, p_all, y_all = 0, 0, [], []
        with torch.no_grad():
            for xg, xz, y in va_dl:
                xg = xg.to(device); xz = xz.to(device); y = y.to(device)
                fg = extract(bb, xg, dtype); fz = extract(bb, xz, dtype)
                logits = head(fg, fz)
                p = F.softmax(logits.float(), dim=-1)
                p_all.append(p.cpu()); y_all.append(y.cpu())
                cs += (logits.argmax(-1) == y).sum().item(); tot += y.size(0)
        acc = cs / max(1, tot)
        if ep % 5 == 0 or ep == args.epochs - 1:
            log(f"DualBranch ep {ep}: loss={tl/max(1,nb):.4f} val_top1={acc:.4f}")
        if acc > best:
            best = acc; best_probs = torch.cat(p_all)
            torch.save({'head_state': head.state_dict(), 'label_to_idx': label_to_idx,
                         'probs': best_probs, 'targets': torch.cat(y_all), 'top1': best,
                         'backbone_arch': args.backbone_arch},
                        args.output)
    log(f"=== DualBranch BEST: {best:.4f}")


# ---------------- 180° rotation consistency ----------------
class RotConsistencyDS(Dataset):
    def __init__(self, manifest, label_to_idx, img_size, is_train, val_fold, min_samples):
        from torchvision import transforms
        cc = Counter()
        self.records = []
        for line in open(manifest):
            r = json.loads(line)
            if r.get('task') != 'classification' or not r.get('unified_label'): continue
            cc[r['unified_label']] += 1
        for line in open(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 not r.get('path') or r.get('storage') != 'fs': continue
            if r.get('integrity_ok') is False: continue
            if cc[r['unified_label']] < min_samples: continue
            fold = r.get('tablet_view_fold', 0)
            if is_train and fold == val_fold: continue
            if not is_train and fold != val_fold: continue
            self.records.append(r)
        self.l2i = label_to_idx
        from torchvision import transforms as T
        self.tf = T.Compose([
            T.Resize((img_size, img_size), antialias=True),
            T.ToTensor(),
            T.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
        ])
        self.is_train = is_train
    def __len__(self): return len(self.records)
    def __getitem__(self, i):
        r = self.records[i]
        img = Image.open(r['path']).convert('RGB')
        t = self.tf(img)
        # 180° rotated view
        t180 = torch.rot90(t, 2, dims=[-2, -1])
        return t, t180, self.l2i[r['unified_label']]


def train_rot_consistency(bb, label_to_idx, manifest, args, device, dtype):
    img_size = get_arch_img_size(args.backbone_arch)
    tr_ds = RotConsistencyDS(manifest, label_to_idx, img_size, True, args.val_fold, args.min_samples)
    va_ds = RotConsistencyDS(manifest, label_to_idx, img_size, False, args.val_fold, args.min_samples)
    log(f"RotConsist Train={len(tr_ds)} Val={len(va_ds)}")
    cc = Counter(tr_ds.l2i[r['unified_label']] for r in tr_ds.records)
    sw = np.array([1.0 / max(1, cc[tr_ds.l2i[r['unified_label']]])**0.5
                    for r in tr_ds.records], dtype=np.float32)
    tr_dl = DataLoader(tr_ds, batch_size=args.batch_size,
                        sampler=WeightedRandomSampler(sw, len(tr_ds), True),
                        num_workers=6, pin_memory=True, drop_last=True)
    va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=4, pin_memory=True)

    with torch.no_grad():
        feat_dim = extract(bb, torch.zeros(1, 3, img_size, img_size, device=device), dtype).size(-1)
    head = nn.Linear(feat_dim, len(label_to_idx)).to(device)
    opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)
    best, best_probs = 0, None
    for ep in range(args.epochs):
        head.train()
        tl, nb = 0, 0
        for x, x180, y in tr_dl:
            x = x.to(device); x180 = x180.to(device); y = y.to(device)
            f1 = extract(bb, x, dtype); f2 = extract(bb, x180, dtype)
            l1 = head(f1); l2 = head(f2)
            # Both predict same class
            ce = 0.5 * (F.cross_entropy(l1, y, label_smoothing=0.1) +
                        F.cross_entropy(l2, y, label_smoothing=0.1))
            # Consistency: KL(softmax(l1) || softmax(l2))
            kl = F.kl_div(F.log_softmax(l1, -1), F.softmax(l2.detach(), -1), reduction='batchmean')
            loss = ce + 0.3 * kl
            opt.zero_grad(); loss.backward(); opt.step()
            tl += loss.item(); nb += 1
        sched.step()
        head.eval()
        cs, tot, p_all, y_all = 0, 0, [], []
        with torch.no_grad():
            for x, _, y in va_dl:
                x = x.to(device); y = y.to(device)
                f = extract(bb, x, dtype); logits = head(f)
                p = F.softmax(logits.float(), -1)
                p_all.append(p.cpu()); y_all.append(y.cpu())
                cs += (logits.argmax(-1) == y).sum().item(); tot += y.size(0)
        acc = cs / max(1, tot)
        if ep % 5 == 0: log(f"RotConsist ep {ep}: loss={tl/max(1,nb):.4f} val={acc:.4f}")
        if acc > best:
            best = acc; best_probs = torch.cat(p_all)
            torch.save({'head_state': head.state_dict(), 'label_to_idx': label_to_idx,
                         'probs': best_probs, 'targets': torch.cat(y_all), 'top1': best,
                         'backbone_arch': args.backbone_arch},
                        args.output)
    log(f"=== RotConsist BEST: {best:.4f}")


# ---------------- MC-Dropout stochastic forward (no train) ----------------
def mc_dropout_eval(bb, label_to_idx, manifest, args, device, dtype):
    """Use existing v4 classifier head + enable dropout at inference, N samples."""
    img_size = get_arch_img_size(args.backbone_arch)
    from torchvision import transforms
    tf = transforms.Compose([
        transforms.Resize((img_size, img_size), antialias=True),
        transforms.ToTensor(),
        transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
    ])
    # Simple val dataset
    recs = []
    for line in open(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
        recs.append(r)

    class VDS(Dataset):
        def __init__(self, rs, tf, l2i): self.r = rs; self.tf = tf; self.l2i = l2i
        def __len__(self): return len(self.r)
        def __getitem__(self, i):
            rr = self.r[i]
            img = Image.open(rr['path']).convert('RGB')
            return self.tf(img), self.l2i[rr['unified_label']]

    dl = DataLoader(VDS(recs, tf, label_to_idx), batch_size=args.batch_size,
                     shuffle=False, num_workers=4, pin_memory=True)
    # Enable dropout: set train mode ONLY on dropout layers
    bb.eval()
    for m in bb.modules():
        if isinstance(m, nn.Dropout): m.train()
    # Collect N stochastic passes
    all_probs, all_y = None, []
    for x, y in dl:
        x = x.to(device, non_blocking=True)
        batch_probs = 0
        N = 8
        for _ in range(N):
            with torch.no_grad(), torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                logits = bb(x)
            batch_probs = batch_probs + F.softmax(logits.float(), -1)
        batch_probs = batch_probs / N
        all_probs = batch_probs.cpu() if all_probs is None else torch.cat([all_probs, batch_probs.cpu()])
        all_y.append(y)
    all_y = torch.cat(all_y)
    acc = (all_probs.argmax(-1) == all_y).float().mean().item()
    log(f"=== MC-Dropout N=8: top1={acc:.4f}")
    torch.save({'probs': all_probs, 'targets': all_y,
                 'label_to_idx': label_to_idx, 'top1': acc, 'N': 8},
                args.output)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--method', required=True,
                    choices=['dual_branch', 'rot_consistency', 'mc_dropout'])
    ap.add_argument('--ckpt', required=True)
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--min-samples', type=int, default=10)
    ap.add_argument('--epochs', type=int, default=25)
    ap.add_argument('--lr', type=float, default=5e-3)
    ap.add_argument('--batch-size', type=int, default=128)
    ap.add_argument('--output', required=True)
    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

    arch, path = args.ckpt.split(':', 1)
    ck = torch.load(path, map_location='cpu', weights_only=False)
    label_to_idx = ck['label_to_idx']
    args.backbone_arch = arch
    img_size = get_arch_img_size(arch)
    bb = build_backbone(arch, n_classes=len(label_to_idx), img_size_override=img_size).to(device)
    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'}
    bb.load_state_dict(sd, strict=False)
    if args.method == 'mc_dropout':
        # For MC we use full model (classifier included)
        mc_dropout_eval(bb, label_to_idx, args.manifest, args, device, dtype)
        return
    for p in bb.parameters(): p.requires_grad = False
    bb.eval()

    if args.method == 'dual_branch':
        train_dual_branch(bb, label_to_idx, args.manifest, args, device, dtype)
    elif args.method == 'rot_consistency':
        train_rot_consistency(bb, label_to_idx, args.manifest, args, device, dtype)

if __name__ == '__main__':
    main()