#!/usr/bin/env python3
"""Decoupled cRT (classifier Re-Training).

Kang et al. ICLR 2020: backbone trained with instance-balanced sampler learns
good representations; tail classes suffer only because head dominates the
classifier. Fix: freeze backbone (with its already-learned features), retrain
ONLY the classifier head with class-balanced sampler.

Often +2-5% top-1 over LDAM+DRW alone.
"""
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 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, HititClsDataset, get_arch_img_size

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

@torch.no_grad()
def extract_features(model, ds, device, dtype, batch_size=128):
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=6)
    raw = model.module if hasattr(model, 'module') else model
    F_all, Y_all = [], []
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        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)
        F_all.append(f.float().cpu()); Y_all.append(y)
    return torch.cat(F_all), torch.cat(Y_all)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True, help='arch:path, trained model')
    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=30)
    ap.add_argument('--lr', type=float, default=1e-2)
    ap.add_argument('--batch-size', type=int, default=256)
    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']
    n_cls = len(label_to_idx)

    model = build_backbone(arch, n_classes=n_cls).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'}
    model.load_state_dict(sd, strict=False)
    for p in model.parameters(): p.requires_grad = False
    model.eval()

    cfg = {'img_size': get_arch_img_size(arch)}
    # Load WITH train split filtering, then override transform to no-aug deterministic
    train_ds = HititClsDataset(args.manifest, cfg, is_train=True,
                                val_fold=args.val_fold, label_to_idx=label_to_idx,
                                min_samples=args.min_samples)
    # Swap transform to eval-style (deterministic resize + norm)
    from torchvision import transforms
    img_size = get_arch_img_size(arch)
    eval_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]),
    ])
    train_ds.tf = eval_tf
    val_ds = HititClsDataset(args.manifest, cfg, is_train=False,
                              val_fold=args.val_fold, label_to_idx=label_to_idx,
                              min_samples=args.min_samples)
    log(f"Train={len(train_ds)} Val={len(val_ds)} classes={n_cls}")

    log("Extracting features (frozen)...")
    tr_f, tr_y = extract_features(model, train_ds, device, dtype, args.batch_size)
    va_f, va_y = extract_features(model, val_ds, device, dtype, args.batch_size)
    D = tr_f.size(-1); log(f"Feat dim: {D}")

    # Class-balanced sampler (inverse freq, NOT sqrt — more aggressive)
    cls_count = Counter(int(y) for y in tr_y)
    sample_w = torch.tensor([1.0 / max(1, cls_count[int(y)]) for y in tr_y], dtype=torch.float32)

    # Simple linear head on features
    head = nn.Linear(D, n_cls).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)

    tr_f_d = tr_f.to(device); tr_y_d = tr_y.to(device)
    va_f_d = va_f.to(device); va_y_d = va_y.to(device)

    best_acc, best_probs = 0.0, None
    N = len(tr_f_d)
    for ep in range(args.epochs):
        head.train()
        # Class-balanced sampling: draw len(tr) indices weighted by sample_w
        idx = torch.multinomial(sample_w, N, replacement=True).to(device)
        for i in range(0, N, 1024):
            b = idx[i:i+1024]
            logits = head(tr_f_d[b])
            loss = F.cross_entropy(logits, tr_y_d[b], label_smoothing=0.1)
            opt.zero_grad(); loss.backward(); opt.step()
        sched.step()
        head.eval()
        with torch.no_grad():
            logits = head(va_f_d)
            acc = (logits.argmax(-1) == va_y_d).float().mean().item()
            _, top5 = logits.topk(5, dim=-1)
            top5_acc = sum(va_y_d[i].item() in top5[i].tolist()
                           for i in range(len(va_y_d))) / len(va_y_d)
        if ep % 5 == 0 or ep == args.epochs - 1:
            log(f"ep {ep}: val_top1={acc:.4f} val_top5={top5_acc:.4f}")
        if acc > best_acc:
            best_acc = acc
            best_probs = F.softmax(head(va_f_d).float(), dim=-1).cpu()

    log(f"=== cRT head retrain: top1 {best_acc:.4f}")

    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    torch.save({'probs': best_probs, 'targets': va_y, 'top1': best_acc,
                 'label_to_idx': label_to_idx, 'head_state': head.state_dict()},
               args.output)
    log(f"Saved → {args.output}")

if __name__ == '__main__':
    main()