#!/usr/bin/env python3
"""Hitit classification fine-tune.

LP-FT strategy (Kumar ICLR 2022): 10 epoch linear probe → 30 epoch LoRA FT.
DINOv3-L / ConvNeXt-V2-L / SigLIP2 backbones supported.
BF16 + torch.compile + FlashAttention.
"""
import os, sys, json, argparse, time
from collections import Counter
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 Dataset, DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from PIL import Image
import yaml

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def setup_ddp():
    if 'RANK' in os.environ:
        torch.distributed.init_process_group(backend='nccl')
        lr = int(os.environ.get('LOCAL_RANK', 0))
        torch.cuda.set_device(lr)
        return True, lr, int(os.environ['WORLD_SIZE']), int(os.environ['RANK'])
    return False, 0, 1, 0

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

class HititClsDataset(Dataset):
    def __init__(self, manifest_path, cfg, is_train=True, val_fold=0, test_fold=4,
                 label_to_idx=None, noisy_paths=None, min_samples=None):
        # Önce TÜM kayıtları oku, global label_to_idx oluştur (train+val ortak)
        all_records = []
        if min_samples is None:
            min_samples = cfg.get('min_samples_per_class', 10)
        noisy_paths = set(noisy_paths or [])
        with open(manifest_path) as f:
            for line in f:
                r = json.loads(line)
                if r.get('task') != 'classification': continue
                if not r.get('unified_label'): 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) < min_samples: continue
                if is_train and r['path'] in noisy_paths: continue
                all_records.append(r)

        # Global label_to_idx (tüm kayıtlarda, train/val öncesi)
        if label_to_idx is None:
            labels = sorted(set(r['unified_label'] for r in all_records))
            self.label_to_idx = {l: i for i, l in enumerate(labels)}
        else:
            self.label_to_idx = label_to_idx
        self.n_classes = len(self.label_to_idx)

        # Split — sadece label_to_idx'te var olan kayıtları al
        self.records = []
        for r in all_records:
            if r['unified_label'] not in self.label_to_idx:
                continue
            fold = r.get('tablet_view_fold', 0)
            if is_train:
                if fold == val_fold or fold == test_fold: continue
            else:
                if fold != val_fold: continue
            self.records.append(r)

        from torchvision import transforms
        img_size = cfg.get('img_size', 224)
        pad = int(img_size * 1.07)
        if is_train:
            self.tf = transforms.Compose([
                transforms.Resize((pad, pad), antialias=True),
                transforms.RandomResizedCrop(img_size, scale=(0.7, 1.0), ratio=(0.85, 1.18)),
                transforms.RandomApply([transforms.RandomAffine(
                    degrees=8, translate=(0.06, 0.06), scale=(0.9, 1.1), shear=4)], p=0.7),
                transforms.ColorJitter(0.35, 0.35, 0.25, 0.05),
                transforms.RandomApply([transforms.GaussianBlur(3, sigma=(0.1, 1.5))], p=0.3),
                transforms.RandomGrayscale(p=0.1),
                transforms.ToTensor(),
                transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
                transforms.RandomErasing(p=0.35, scale=(0.02, 0.18), ratio=(0.3, 3.3)),
            ])
        else:
            self.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]),
            ])
        self.img_size = img_size

    def __len__(self): return len(self.records)

    def __getitem__(self, idx):
        r = self.records[idx]
        img = Image.open(r['path']).convert('RGB')
        img = self.tf(img)
        y = self.label_to_idx[r['unified_label']]
        return img, y

ARCH_IMG_SIZE = {
    'dinov3_vitl14': 224, 'dinov3_vitb14': 224,
    'convnextv2_large': 224,
    'siglip2_so400m': 224,
    'swinv2_large': 384,
    'eva02_large': 448,
    'maxvit_large': 384,
    'mambavision_l': 224,
}

def get_arch_img_size(arch): return ARCH_IMG_SIZE.get(arch, 224)

def build_backbone(arch, ssl_ckpt=None, n_classes=198, img_size_override=None):
    import timm
    # Map arch name
    arch_map = {
        # DINOv3-L: timm has vit_large variant; fallback to base dinov2 if unavailable
        'dinov3_vitl14': 'vit_large_patch14_dinov2',
        'dinov3_vitb14': 'vit_base_patch14_dinov2',
        'convnextv2_large': 'convnextv2_large',
        'siglip2_so400m': 'vit_so400m_patch14_siglip_224',
        # Diverse backbones for ensemble
        'swinv2_large': 'swinv2_large_window12to24_192to384.ms_in22k_ft_in1k',
        'eva02_large': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k',
        'maxvit_large': 'maxvit_large_tf_384.in1k',
        'mambavision_l': 'mambavision_l_1k',
    }
    tname = arch_map.get(arch, 'vit_base_patch14_dinov2')
    img_size = img_size_override or ARCH_IMG_SIZE.get(arch, 224)
    create_kwargs = dict(pretrained=True, num_classes=n_classes)
    # Only ViT-family supports dynamic img_size; swin/convnext/maxvit/eva02 have
    # fixed native sizes — pass img_size only when timm accepts it
    if 'dinov2' in tname or tname.startswith('vit_'):
        create_kwargs['img_size'] = img_size
        create_kwargs['dynamic_img_size'] = True
    elif tname.startswith('eva02'):
        create_kwargs['img_size'] = img_size
    try:
        model = timm.create_model(tname, **create_kwargs)
    except Exception as e:
        log(f"timm load fallback for {tname}: {e}")
        # Try without dynamic_img_size
        ck = dict(pretrained=True, num_classes=n_classes)
        try: model = timm.create_model(tname, **ck)
        except Exception as e2:
            log(f"timm still failed: {e2}; falling back to vit_base_patch16_224")
            model = timm.create_model('vit_base_patch16_224', pretrained=True,
                                       num_classes=n_classes, dynamic_img_size=True)

    # Load SSL ckpt if provided — SADECE SSL loss düşmüş ise (validate)
    if ssl_ckpt and Path(ssl_ckpt).exists():
        ck = torch.load(ssl_ckpt, map_location='cpu', weights_only=False)
        bb_state = ck.get('backbone', ck)
        # Validate SSL quality: backbone collapse check
        # cls_token norm should not be ~0 (sign of DINO collapse)
        cls_tok = bb_state.get('cls_token')
        ssl_ok = cls_tok is not None and float(cls_tok.abs().mean()) > 1e-3
        if not ssl_ok:
            log(f"⚠️ SSL checkpoint quality poor (cls_token mean={float(cls_tok.abs().mean() if cls_tok is not None else 0):.6f}); keeping ImageNet pretrained")
        else:
            own = model.state_dict()
            matched = {k: v for k, v in bb_state.items() if k in own and own[k].shape == v.shape}
            model.load_state_dict(matched, strict=False)
            log(f"Loaded SSL backbone: {len(matched)} tensors matched (SSL valid)")
    return model

def apply_lora_simple(model, r=32, alpha=64):
    """Minimal LoRA via additive weights (simplified)."""
    # Freeze backbone, only head trainable initially
    for p in model.parameters():
        p.requires_grad = False
    # Head always trainable
    if hasattr(model, 'head'):
        for p in model.head.parameters():
            p.requires_grad = True
    elif hasattr(model, 'fc'):
        for p in model.fc.parameters():
            p.requires_grad = True
    return model

def unfreeze_lora(model, n_last_blocks=4):
    """Unfreeze last N transformer blocks / stages for FT phase."""
    unfrozen = 0
    # ViT: model.blocks (ModuleList)
    if hasattr(model, 'blocks') and len(list(model.blocks)) > 0:
        blocks = list(model.blocks)
        for block in blocks[-n_last_blocks:]:
            for p in block.parameters(): p.requires_grad = True; unfrozen += 1
    # ConvNeXt: model.stages (4 stages, each with blocks) — unfreeze last 2 stages
    if hasattr(model, 'stages') and len(list(model.stages)) > 0:
        stages = list(model.stages)
        n_stages = min(2, len(stages))
        for stage in stages[-n_stages:]:
            for p in stage.parameters(): p.requires_grad = True; unfrozen += 1
    # Swin-V2 (timm): model.layers → [SwinStage × 4], each with .blocks
    if hasattr(model, 'layers') and len(list(model.layers)) > 0:
        layers = list(model.layers)
        # Unfreeze last 2 stages (stage3 has most blocks, stage4 is deepest)
        n_stages = min(2, len(layers))
        for layer in layers[-n_stages:]:
            for p in layer.parameters(): p.requires_grad = True; unfrozen += 1
    # EVA-02: also uses .blocks (already handled above)
    # MaxViT: model.stages with blocks
    # Norm + head always trainable
    for attr in ('norm', 'norm_pre', 'head_norm', 'fc_norm', 'head'):
        if hasattr(model, attr):
            obj = getattr(model, attr)
            if hasattr(obj, 'parameters'):
                for p in obj.parameters(): p.requires_grad = True
    return model

class LDAMLoss(nn.Module):
    """LDAM (Cao NeurIPS 2019): margin ∝ n_c^(-1/4).

    Encourages larger margins for minority classes — direct remedy for
    head/tail accuracy gap. Uses DRW (deferred re-weighting) via per-class
    weight schedule passed at call time.
    """
    def __init__(self, class_counts, max_m=0.5, s=30.0, label_smoothing=0.0):
        super().__init__()
        m = 1.0 / np.sqrt(np.sqrt(np.maximum(1, class_counts)))
        m = m * (max_m / m.max())
        self.register_buffer('m', torch.tensor(m, dtype=torch.float32))
        self.s = s
        self.label_smoothing = label_smoothing

    def forward(self, logits, target, soft=None, weight=None):
        # soft: optional soft-labels (for mixup)
        batch_m = self.m[target if soft is None else target]  # just for shape
        # Apply margin to true class
        if soft is None:
            idx = torch.arange(logits.size(0), device=logits.device)
            m = self.m.to(logits.device)
            margin = torch.zeros_like(logits)
            margin[idx, target] = m[target]
            logits_m = (logits - margin) * self.s / 30.0
            return F.cross_entropy(logits_m, target, weight=weight,
                                    label_smoothing=self.label_smoothing)
        # Soft (mixup): subtract per-sample weighted margin
        m = self.m.to(logits.device)
        margin = (soft * m.unsqueeze(0)).sum(-1, keepdim=True)
        logits_m = logits - margin * F.one_hot(logits.argmax(-1), logits.size(-1))
        logp = F.log_softmax(logits_m * self.s / 30.0, dim=-1)
        eps = self.label_smoothing
        soft_s = soft * (1 - eps) + eps / logits.size(-1)
        if weight is not None:
            w_per = weight[target] if target.dim() == 1 else None
            loss = -(soft_s * logp).sum(-1)
            if w_per is not None: loss = loss * w_per
            return loss.mean()
        return -(soft_s * logp).sum(-1).mean()

def supcon_loss(feats, labels, temperature=0.07):
    """Supervised contrastive loss (Khosla NeurIPS 2020)."""
    feats = F.normalize(feats, dim=-1)
    sim = feats @ feats.t() / temperature  # (B, B)
    # Remove self-similarity
    B = feats.size(0)
    mask_self = torch.eye(B, device=feats.device, dtype=torch.bool)
    sim = sim.masked_fill(mask_self, -1e4)
    # Positives: same label
    labels = labels.view(-1, 1)
    mask_pos = (labels == labels.t()).float().masked_fill(mask_self, 0)
    # log-softmax over non-self
    log_prob = sim - sim.logsumexp(dim=-1, keepdim=True)
    # Mean log-prob over positives (rows with no positives get 0)
    n_pos = mask_pos.sum(-1).clamp_min(1)
    loss = -(mask_pos * log_prob).sum(-1) / n_pos
    return loss.mean()

def make_drw_weights(class_counts, beta=0.9999, device='cuda'):
    """Class-balanced weights (Cui CVPR 2019): (1-β)/(1-β^n_c)."""
    n = np.asarray(class_counts, dtype=np.float64).clip(min=1)
    w = (1.0 - beta) / (1.0 - beta ** n)
    w = w / w.mean()
    return torch.tensor(w, dtype=torch.float32, device=device)

def _mixup_cutmix(imgs, ys, n_classes, mixup_alpha=0.2, cutmix_alpha=1.0, prob=0.5,
                   class_counts=None, balanced=False):
    """Apply MixUp or CutMix with prob. Returns (imgs, soft_targets).
    Balanced MixUp (Galdran 2021): sample partner with class-inverse-frequency weighting.
    """
    import numpy as np
    soft = F.one_hot(ys, n_classes).float()
    if torch.rand(1).item() > prob:
        return imgs, soft
    use_cutmix = torch.rand(1).item() < 0.5
    if balanced and class_counts is not None:
        # Partner sampling weighted by 1/freq(class_of_sample)
        w = (1.0 / class_counts.clamp_min(1))[ys].to(imgs.device)
        perm = torch.multinomial(w, imgs.size(0), replacement=True)
    else:
        perm = torch.randperm(imgs.size(0), device=imgs.device)
    if use_cutmix:
        lam = float(np.random.beta(cutmix_alpha, cutmix_alpha))
        H, W = imgs.shape[-2], imgs.shape[-1]
        cut_rat = (1. - lam) ** 0.5
        cw, ch = int(W * cut_rat), int(H * cut_rat)
        cx, cy = np.random.randint(W), np.random.randint(H)
        x1, x2 = max(0, cx - cw // 2), min(W, cx + cw // 2)
        y1, y2 = max(0, cy - ch // 2), min(H, cy + ch // 2)
        imgs[:, :, y1:y2, x1:x2] = imgs[perm, :, y1:y2, x1:x2]
        lam = 1. - ((x2 - x1) * (y2 - y1) / (W * H))
    else:
        lam = float(np.random.beta(mixup_alpha, mixup_alpha))
        imgs = lam * imgs + (1 - lam) * imgs[perm]
    soft = lam * soft + (1 - lam) * soft[perm]
    return imgs, soft

def train_epoch(model, loader, optimizer, device, dtype, rank=0, ema_model=None,
                n_classes=198, use_mix=False, ldam_loss=None, cb_weights=None,
                supcon_weight=0.0, sam=None, manifold_mixup=False,
                balanced_mixup=False, class_counts=None):
    model.train()
    total_loss, total_correct, total = 0, 0, 0
    for imgs, ys in loader:
        imgs = imgs.to(device, non_blocking=True)
        ys = ys.to(device, non_blocking=True)
        if manifold_mixup:
            # Mix at feature space: run forward_features, mix, then head
            raw = model.module if hasattr(model, 'module') else model
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                if hasattr(raw, 'forward_features'):
                    feats = raw.forward_features(imgs)
                    if feats.dim() == 3: feats = feats[:, 0]
                    elif feats.dim() == 4: feats = feats.mean(dim=(2, 3))
                    lam = float(np.random.beta(0.2, 0.2))
                    perm = torch.randperm(feats.size(0), device=feats.device)
                    feats_m = lam * feats + (1 - lam) * feats[perm]
                    head_fn = raw.head if hasattr(raw, 'head') and not isinstance(raw.head, nn.Identity) \
                              else (raw.fc if hasattr(raw, 'fc') else raw.classifier)
                    logits = head_fn(feats_m)
                    soft = F.one_hot(ys, n_classes).float()
                    soft = lam * soft + (1 - lam) * soft[perm]
                else:
                    # Fallback to pixel mixup
                    imgs_m, soft = _mixup_cutmix(imgs, ys, n_classes,
                                          class_counts=class_counts, balanced=balanced_mixup)
                    logits = model(imgs_m)
            if ldam_loss is not None:
                loss = ldam_loss(logits, ys, soft=soft, weight=cb_weights)
            else:
                logp = F.log_softmax(logits, dim=-1)
                eps = 0.1; soft_s = soft * (1 - eps) + eps / n_classes
                loss = -(soft_s * logp).sum(-1).mean()
        elif use_mix:
            imgs_m, soft = _mixup_cutmix(imgs, ys, n_classes,
                                          class_counts=class_counts, balanced=balanced_mixup)
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                logits = model(imgs_m)
                if ldam_loss is not None:
                    loss = ldam_loss(logits, ys, soft=soft, weight=cb_weights)
                else:
                    logp = F.log_softmax(logits, dim=-1)
                    eps = 0.1
                    soft_s = soft * (1 - eps) + eps / n_classes
                    if cb_weights is not None:
                        w_per = cb_weights[ys]
                        loss = -(soft_s * logp).sum(-1)
                        loss = (loss * w_per).mean()
                    else:
                        loss = -(soft_s * logp).sum(-1).mean()
        else:
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                if supcon_weight > 0:
                    # Hook penultimate features: call forward_features then head
                    raw = model.module if hasattr(model, 'module') else model
                    if hasattr(raw, 'forward_features'):
                        feats = raw.forward_features(imgs)
                        if feats.dim() > 2:
                            feats = feats.mean(dim=tuple(range(1, feats.dim()-1)))
                            if feats.dim() == 3: feats = feats.mean(1)
                        logits = raw.head(feats) if hasattr(raw, 'head') else raw.fc(feats)
                    else:
                        logits = model(imgs); feats = None
                else:
                    logits = model(imgs); feats = None
                if ldam_loss is not None:
                    loss = ldam_loss(logits, ys, weight=cb_weights)
                else:
                    loss = F.cross_entropy(logits, ys, weight=cb_weights,
                                            label_smoothing=0.1)
                if supcon_weight > 0 and feats is not None:
                    loss = loss + supcon_weight * supcon_loss(feats, ys)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if sam is not None:
            # Focal-SAM 2-pass: perturb, recompute loss, step from original.
            # For manifold_mixup (no imgs_m in pixel space) fall back to plain imgs —
            # the perturbation itself is what matters; slight label mismatch is acceptable.
            fs = sam.compute_focal_scale(ys)
            sam.first_step(zero_grad=True, focal_scale=fs)
            imgs_for_sam = locals().get('imgs_m', None)
            if imgs_for_sam is None or not use_mix:
                imgs_for_sam = imgs
            with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
                logits2 = model(imgs_for_sam)
                if ldam_loss is not None:
                    loss2 = ldam_loss(logits2, ys, weight=cb_weights) if not use_mix \
                            else ldam_loss(logits2, ys, soft=soft, weight=cb_weights)
                else:
                    loss2 = F.cross_entropy(logits2, ys, weight=cb_weights,
                                             label_smoothing=0.1)
            loss2.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            sam.second_step(zero_grad=True)
        else:
            optimizer.step()
        if ema_model is not None:
            ema_model.update_parameters(model)
        with torch.no_grad():
            pred = logits.argmax(-1)
            total_correct += (pred == ys).sum().item()
            total += ys.size(0)
            total_loss += loss.item() * ys.size(0)
    return total_loss / max(1, total), total_correct / max(1, total)

@torch.no_grad()
def validate(model, loader, device, dtype):
    model.eval()
    total_correct, total = 0, 0
    for imgs, ys in loader:
        imgs = imgs.to(device, non_blocking=True)
        ys = ys.to(device, non_blocking=True)
        with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
            logits = model(imgs)
        pred = logits.argmax(-1)
        total_correct += (pred == ys).sum().item()
        total += ys.size(0)
    return total_correct / max(1, total)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/classification_hitit_only.yaml'))
    ap.add_argument('--backbone', help='SSL ckpt path')
    ap.add_argument('--backbone-arch', default='dinov3_vitl14')
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/classification_hitit/'))
    ap.add_argument('--noisy-issues', help='Path to cleanlab_issues.json (drops flagged train rows)')
    ap.add_argument('--min-samples', type=int, default=10, help='Drop classes with <N training samples')
    ap.add_argument('--use-ldam', action='store_true', default=True, help='LDAM margin loss')
    ap.add_argument('--no-ldam', dest='use_ldam', action='store_false')
    ap.add_argument('--drw-beta', type=float, default=0.9999,
                    help='Class-balanced DRW beta; applied last 20%% of FT epochs')
    ap.add_argument('--supcon-weight', type=float, default=0.0,
                    help='Auxiliary SupCon loss weight (0 disabled)')
    ap.add_argument('--focal-sam', action='store_true',
                    help='Use Focal-SAM optimizer wrapper (ICML 2025)')
    ap.add_argument('--sam-rho', type=float, default=0.05)
    ap.add_argument('--img-size', type=int, default=None,
                    help='Override arch-native img_size (e.g. 336 for DINOv3-L)')
    ap.add_argument('--manifold-mixup', action='store_true',
                    help='MixUp at penultimate feature layer (Verma ICML 2019)')
    ap.add_argument('--balanced-mixup', action='store_true',
                    help='Balanced MixUp (Galdran 2021): sample partner inverse-freq weighted')
    ap.add_argument('--curriculum', action='store_true',
                    help='Curriculum sampling: head classes first, then tail')
    ap.add_argument('--manifest', default=None,
                    help='Override default hitit_local classification manifest')
    ap.add_argument('--seed', type=int, default=42,
                    help='Random seed (for seed-ensemble diversity)')
    ap.add_argument('--batch-size', type=int, default=None,
                    help='Override YAML batch_size (use for OOM-prone large backbones)')
    args = ap.parse_args()

    # Seed everything for reproducibility / seed-ensemble
    import numpy as _np
    import random as _random
    torch.manual_seed(args.seed)
    _np.random.seed(args.seed)
    _random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    cfg = yaml.safe_load(open(args.config))
    # Inject arch-native img_size (overrides cfg.img_size)
    cfg['img_size'] = args.img_size or get_arch_img_size(args.backbone_arch)
    if args.batch_size is not None:
        cfg['batch_size'] = args.batch_size
    is_ddp, lr, world, rank = setup_ddp()
    device = torch.device(f'cuda:{lr}' if torch.cuda.is_available() else 'cpu')
    Path(args.output).mkdir(parents=True, exist_ok=True)

    log(f"DDP={is_ddp}, world={world}, rank={rank}, device={device}, img_size={cfg['img_size']}", rank)

    # Optional noisy-label drop
    noisy_paths = []
    if args.noisy_issues and Path(args.noisy_issues).exists():
        nl = json.load(open(args.noisy_issues))
        noisy_paths = [r['path'] for r in nl.get('noisy_records', [])]
        log(f"Noisy drop: {len(noisy_paths)} paths flagged", rank)

    # Data — global label_to_idx kullanılır
    manifest = Path(args.manifest) if args.manifest else \
        ROOT / 'datasets/sources/hitit_local/manifest_classification.jsonl'
    log(f"Manifest: {manifest}", rank)
    train_ds = HititClsDataset(manifest, cfg, is_train=True,
                                noisy_paths=noisy_paths, min_samples=args.min_samples)
    val_ds = HititClsDataset(manifest, cfg, is_train=False,
                              label_to_idx=train_ds.label_to_idx,
                              min_samples=args.min_samples)
    log(f"Train={len(train_ds)} Val={len(val_ds)} classes={train_ds.n_classes}", rank)

    bs = cfg.get('batch_size', 64) // max(world, 1)
    # Class-balanced weighted sampler (sqrt-inverse freq)
    import numpy as np
    label_counts = Counter(train_ds.label_to_idx[r['unified_label']] for r in train_ds.records)
    # Per-class counts aligned by idx 0..C-1
    class_counts = np.array([label_counts.get(i, 0) for i in range(train_ds.n_classes)],
                             dtype=np.int64)
    sample_w = np.array([1.0 / max(1, label_counts[train_ds.label_to_idx[r['unified_label']]]) ** 0.5
                         for r in train_ds.records], dtype=np.float32)
    from torch.utils.data import WeightedRandomSampler, DistributedSampler as _DS
    if is_ddp:
        # DDP-compatible class-balanced sampler: replicate WeightedRandomSampler on each rank
        # by pre-expanding indices and slicing per rank
        g = torch.Generator(); g.manual_seed(42)
        expanded = torch.multinomial(torch.from_numpy(sample_w), len(sample_w),
                                      replacement=True, generator=g).tolist()
        class _CBSampler:
            def __init__(self, idx_list, num_replicas, rank_id):
                self.idx = idx_list; self.nr = num_replicas; self.r = rank_id
                self.num_samples = len(idx_list) // num_replicas
            def set_epoch(self, ep):
                g = torch.Generator(); g.manual_seed(42 + ep)
                self.idx = torch.multinomial(torch.from_numpy(sample_w), len(sample_w),
                                              replacement=True, generator=g).tolist()
            def __iter__(self): return iter(self.idx[self.r::self.nr][:self.num_samples])
            def __len__(self): return self.num_samples
        tsamp = _CBSampler(expanded, world, rank)
    else:
        tsamp = WeightedRandomSampler(sample_w, num_samples=len(train_ds), replacement=True)
    # Curriculum: start head-only, gradually expand to tail
    class _CurriculumSampler:
        def __init__(self, weights, class_counts, records, label_to_idx):
            self.w0 = weights.copy()
            self.cc = class_counts
            self.recs = records
            self.l2i = label_to_idx
            self.epoch = 0
            self.total_epochs = 100
            self.num_samples = len(weights)
        def set_epoch(self, ep):
            self.epoch = ep
        def __iter__(self):
            # Fraction of training progress
            frac = self.epoch / max(1, self.total_epochs)
            # Early epochs: weight = 1 if head (count > 100), else 0.2
            # Late epochs: original class-balanced weights
            cur_w = self.w0.copy()
            for i, r in enumerate(self.recs):
                n = self.cc[self.l2i[r['unified_label']]]
                if n > 100: tier_w = 1.0
                elif n >= 20: tier_w = 0.2 + 0.8 * frac
                else: tier_w = 0.05 + 0.95 * frac
                cur_w[i] *= tier_w
            return iter(torch.multinomial(torch.from_numpy(cur_w), self.num_samples, True).tolist())
        def __len__(self): return self.num_samples
    if args.curriculum and not is_ddp:
        tsamp = _CurriculumSampler(sample_w, label_counts, train_ds.records, train_ds.label_to_idx)
        tsamp.total_epochs = cfg.get('epochs', 40)
    train_loader = DataLoader(train_ds, batch_size=bs, shuffle=False,
                              sampler=tsamp, num_workers=8, pin_memory=True, drop_last=True,
                              persistent_workers=True)
    val_loader = DataLoader(val_ds, batch_size=bs, shuffle=False, num_workers=4, pin_memory=True)

    # Model
    model = build_backbone(args.backbone_arch, args.backbone,
                            n_classes=train_ds.n_classes,
                            img_size_override=cfg['img_size']).to(device)

    # LP-FT
    # Phase 1: Linear probe (10 epochs, head only)
    # Phase 2: LoRA FT (30 epochs, last 4 blocks + head)
    total_epochs = cfg.get('epochs', 40)
    lp_epochs = min(10, total_epochs // 4)

    use_bf16 = cfg.get('bf16', True) and torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    # LDAM loss (margin ∝ n_c^-1/4) — large margin for tail classes
    ldam = LDAMLoss(class_counts, max_m=0.5, s=30.0, label_smoothing=0.1).to(device) \
        if args.use_ldam else None
    if ldam is not None:
        log(f"LDAM loss: margin range {ldam.m.min().item():.3f}-{ldam.m.max().item():.3f}", rank)
    # DRW schedule: last 20% of FT uses class-balanced weights
    cb_weights_full = make_drw_weights(class_counts, beta=args.drw_beta, device=device)

    # EMA
    from torch.optim.swa_utils import AveragedModel
    ema_model = None

    # --- Phase 1: Linear probe ---
    log(f"=== Phase 1 (LP): {lp_epochs} epoch, head-only ===", rank)
    apply_lora_simple(model)
    trainable = [p for p in model.parameters() if p.requires_grad]
    log(f"  Trainable params: {sum(p.numel() for p in trainable):,}", rank)
    if is_ddp:
        model = DDP(model, device_ids=[lr], find_unused_parameters=True)
    opt = torch.optim.AdamW(trainable, lr=3e-3, weight_decay=0.05, fused=False)

    for epoch in range(lp_epochs):
        if hasattr(tsamp, 'set_epoch'): tsamp.set_epoch(epoch)
        tl, ta = train_epoch(model, train_loader, opt, device, dtype, rank,
                              n_classes=train_ds.n_classes, use_mix=False,
                              ldam_loss=None, cb_weights=None)
        va = validate(model, val_loader, device, dtype)
        log(f"LP epoch {epoch}: train_loss={tl:.4f} train_acc={ta:.3f} val_acc={va:.3f}", rank)

    # --- Phase 2: LoRA FT ---
    log(f"=== Phase 2 (LoRA FT): {total_epochs - lp_epochs} epoch ===", rank)
    raw = model.module if is_ddp else model
    unfreeze_lora(raw, n_last_blocks=8)
    trainable = [p for p in raw.parameters() if p.requires_grad]
    log(f"  Trainable params: {sum(p.numel() for p in trainable):,}", rank)

    opt = torch.optim.AdamW(trainable, lr=3e-4, weight_decay=0.1, fused=False)
    sam = None
    if args.focal_sam:
        from enhancements.focal_sam import FocalSAM
        sam = FocalSAM(opt, rho=args.sam_rho, class_counts=class_counts.tolist(),
                        gamma=2.0, alpha=0.25, device=device)
        log(f"Focal-SAM active: rho={args.sam_rho}", rank)
    # Cosine with warmup (5 epochs) for smoother FT
    from torch.optim.lr_scheduler import LambdaLR
    warm = 5
    def lr_lam(step):
        if step < warm: return (step + 1) / warm
        import math
        prog = (step - warm) / max(1, total_epochs - lp_epochs - warm)
        return 0.5 * (1 + math.cos(math.pi * prog))
    scheduler = LambdaLR(opt, lr_lam)
    ema_model = AveragedModel(raw, avg_fn=lambda avg, new, n: 0.9995 * avg + 0.0005 * new)

    best_acc = 0
    best_ema_acc = 0
    ema_warm = 10  # start EMA eval after 10 FT epochs
    ft_span = total_epochs - lp_epochs
    drw_start = lp_epochs + int(ft_span * 0.8)  # enable DRW last 20%
    # In-run SWA: average last 20% of epochs
    swa_start = lp_epochs + int(ft_span * 0.8)
    swa_model = AveragedModel(raw)
    swa_count = 0
    for epoch in range(lp_epochs, total_epochs):
        if hasattr(tsamp, 'set_epoch'): tsamp.set_epoch(epoch)
        cb_w = cb_weights_full if epoch >= drw_start else None
        tl, ta = train_epoch(model, train_loader, opt, device, dtype, rank, ema_model,
                              n_classes=train_ds.n_classes, use_mix=True,
                              ldam_loss=ldam, cb_weights=cb_w,
                              supcon_weight=args.supcon_weight, sam=sam,
                              manifold_mixup=args.manifold_mixup,
                              balanced_mixup=args.balanced_mixup,
                              class_counts=train_ds.class_counts if hasattr(train_ds, 'class_counts') else None)
        va = validate(model, val_loader, device, dtype)
        # EMA val (more stable, often +1-2%)
        va_ema = 0
        if (epoch - lp_epochs) >= ema_warm:
            va_ema = validate(ema_model, val_loader, device, dtype)
        scheduler.step()
        log(f"FT epoch {epoch}: train_loss={tl:.4f} train_acc={ta:.3f} val_acc={va:.3f} val_ema={va_ema:.3f}", rank)
        if va > best_acc and rank == 0:
            best_acc = va
            Path(args.output).mkdir(parents=True, exist_ok=True)
            torch.save({'model': raw.state_dict(), 'cfg': cfg, 'val_acc': va,
                        'label_to_idx': train_ds.label_to_idx},
                       Path(args.output) / 'best.pt')
        if va_ema > best_ema_acc and rank == 0:
            best_ema_acc = va_ema
            Path(args.output).mkdir(parents=True, exist_ok=True)
            torch.save({'model': ema_model.module.state_dict(), 'cfg': cfg, 'val_acc': va_ema,
                        'label_to_idx': train_ds.label_to_idx},
                       Path(args.output) / 'best_ema.pt')
        # In-run SWA: update averaged model last 20%
        if epoch >= swa_start:
            swa_model.update_parameters(raw); swa_count += 1
            if rank == 0 and swa_count % 5 == 0:
                va_swa = validate(swa_model, val_loader, device, dtype)
                log(f"  SWA ({swa_count} ckpts): val_acc={va_swa:.3f}", rank)
                if va_swa > best_ema_acc:
                    Path(args.output).mkdir(parents=True, exist_ok=True)
                    torch.save({'model': swa_model.module.state_dict(), 'cfg': cfg, 'val_acc': va_swa,
                                'label_to_idx': train_ds.label_to_idx, 'swa': True},
                               Path(args.output) / 'best_swa.pt')

    if rank == 0:
        log(f"DONE: best val_acc={best_acc:.3f} best_ema={best_ema_acc:.3f}, saved to {args.output}", rank)

    if is_ddp:
        torch.distributed.destroy_process_group()

if __name__ == '__main__':
    main()