#!/usr/bin/env python3
"""TransFG-style part-level ViT head.

Uses CLS→patch attention from last block to select top-K informative patches,
pools their tokens, and classifies. Captures wedge-level detail — Hittite
signs differ by stroke configuration, not global appearance.

Reference: He ECCV 2022 "TransFG: A Transformer Architecture for Fine-grained Recognition"
"""
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

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, HititClsDataset

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


@torch.no_grad()
def extract_tokens_and_cls_attn(raw_model, x, dtype):
    """Extract patch tokens + CLS-to-patch attention from last ViT block."""
    # Timm ViT: forward_features returns (B, N, D) with CLS at index 0
    attn_weights = {}
    last_block = raw_model.blocks[-1]
    orig_forward = last_block.attn.forward

    def capturing_forward(xin, **kwargs):
        # timm ≥1.0 passes attn_mask/is_causal; we ignore them and use the
        # monkey-patched self-attention path (no masking needed for ViT CLS).
        B, N, _ = xin.shape
        qkv = last_block.attn.qkv(xin).reshape(B, N, 3, last_block.attn.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        a = (q @ k.transpose(-2, -1)) * last_block.attn.scale
        a = a.softmax(-1)
        attn_weights['a'] = a.mean(1)  # (B, N, N) head-averaged
        out = (a @ v).transpose(1, 2).reshape(B, N, -1)
        out = last_block.attn.proj(out)
        return last_block.attn.proj_drop(out)

    last_block.attn.forward = capturing_forward
    with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
        feats = raw_model.forward_features(x)
    last_block.attn.forward = orig_forward

    # feats: (B, N, D), attn: (B, N, N); CLS row = attn[:, 0, 1:]
    cls_attn = attn_weights['a'][:, 0, 1:]  # attention from CLS to patches
    patch_tokens = feats[:, 1:]  # (B, P, D)
    cls_token = feats[:, 0]      # (B, D)
    return cls_token.float(), patch_tokens.float(), cls_attn.float()


class PartViTHead(nn.Module):
    def __init__(self, feat_dim, n_classes, k_parts=8):
        super().__init__()
        self.k = k_parts
        # CLS-informed classifier + part pooling
        self.cls_fc = nn.Linear(feat_dim, n_classes)
        self.part_fc = nn.Linear(feat_dim, n_classes)
        # Learnable attention refinement over selected parts
        self.part_attn = nn.Sequential(
            nn.Linear(feat_dim, feat_dim // 4), nn.GELU(),
            nn.Linear(feat_dim // 4, 1)
        )

    def forward(self, cls_tok, patch_tokens, cls_attn):
        B, P, D = patch_tokens.shape
        # Pick top-K patches by CLS attention
        topk = cls_attn.topk(self.k, dim=-1).indices  # (B, k)
        picked = torch.gather(patch_tokens, 1, topk.unsqueeze(-1).expand(-1, -1, D))  # (B, k, D)
        # Learnable attention over picked tokens
        a = self.part_attn(picked).softmax(1)  # (B, k, 1)
        part_feat = (picked * a).sum(1)  # (B, D)
        # Two-head: CLS + part, averaged
        return 0.5 * self.cls_fc(cls_tok) + 0.5 * self.part_fc(part_feat)


def main():
    ap = argparse.ArgumentParser()
    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('--k-parts', type=int, default=8)
    ap.add_argument('--batch-size', type=int, default=64)
    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)
    img_size = get_arch_img_size(arch)
    backbone = build_backbone(arch, n_classes=n_cls, 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'}
    backbone.load_state_dict(sd, strict=False)
    for p in backbone.parameters(): p.requires_grad = False
    backbone.eval()
    raw = backbone.module if hasattr(backbone, 'module') else backbone

    cfg = {'img_size': img_size}
    tr_ds = HititClsDataset(args.manifest, cfg, is_train=True,
                             val_fold=args.val_fold, label_to_idx=label_to_idx,
                             min_samples=args.min_samples)
    va_ds = HititClsDataset(args.manifest, cfg, is_train=False,
                             val_fold=args.val_fold, label_to_idx=label_to_idx,
                             min_samples=args.min_samples)
    cc = Counter(tr_ds.label_to_idx[r['unified_label']] for r in tr_ds.records)
    sample_w = np.array([1.0 / max(1, cc[tr_ds.label_to_idx[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(sample_w, 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)

    # Get feat_dim from dummy
    with torch.no_grad():
        dummy = torch.zeros(1, 3, img_size, img_size, device=device)
        cls_t, tok, att = extract_tokens_and_cls_attn(raw, dummy, dtype)
        feat_dim = cls_t.size(-1)
    log(f"Feat dim {feat_dim}, n_patches {tok.size(1)}")

    head = PartViTHead(feat_dim, n_cls, k_parts=args.k_parts).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, best_y = 0, None, None
    for ep in range(args.epochs):
        head.train()
        tl, nb = 0, 0
        for x, y in tr_dl:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            cls_t, tok, att = extract_tokens_and_cls_attn(raw, x, dtype)
            logits = head(cls_t, tok, att)
            loss = F.cross_entropy(logits, y, label_smoothing=0.1)
            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, non_blocking=True); y = y.to(device, non_blocking=True)
                cls_t, tok, att = extract_tokens_and_cls_attn(raw, x, dtype)
                logits = head(cls_t, tok, att)
                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 or ep == args.epochs - 1:
            log(f"ep {ep}: loss={tl/max(1,nb):.4f} val_top1={acc:.4f}")
        if acc > best:
            best = acc
            best_probs = torch.cat(p_all); best_y = torch.cat(y_all)
            torch.save({'head_state': head.state_dict(), 'label_to_idx': label_to_idx,
                         'probs': best_probs, 'targets': best_y, 'top1': best,
                         'backbone_arch': arch, 'k_parts': args.k_parts},
                        args.output)
    log(f"=== Part-ViT BEST: {best:.4f}")

if __name__ == '__main__':
    main()