#!/usr/bin/env python3
"""Multi-task: classification + bbox-normalized position (reading order hint).

Auxiliary head predicts (x_center, y_center) normalized to tablet bbox.
Forces backbone to encode spatial context → better discrimination for
signs that only differ in tablet position (numerals, logograms).
"""
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)


class PosDS(Dataset):
    def __init__(self, manifest, label_to_idx, img_size, is_train, val_fold, min_samples):
        cls_count = Counter()
        self.records = []
        for line in open(manifest):
            r = json.loads(line)
            if r.get('task') != 'classification' or not r.get('unified_label'): continue
            cls_count[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 cls_count[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
        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]),
        ])
    def __len__(self): return len(self.records)
    def __getitem__(self, i):
        r = self.records[i]
        img = Image.open(r['path']).convert('RGB')
        # position from bbox if available, else 0.5,0.5
        bbox = r.get('bbox_normalized') or r.get('bbox') or [0.25, 0.25, 0.75, 0.75]
        if isinstance(bbox, list) and len(bbox) == 4:
            xc = 0.5 * (bbox[0] + bbox[2]); yc = 0.5 * (bbox[1] + bbox[3])
        else:
            xc, yc = 0.5, 0.5
        pos = torch.tensor([xc, yc], dtype=torch.float32)
        return self.tf(img), self.l2i[r['unified_label']], pos


class MultiTaskHead(nn.Module):
    def __init__(self, feat_dim, n_classes):
        super().__init__()
        self.cls = nn.Linear(feat_dim, n_classes)
        self.pos = nn.Sequential(
            nn.Linear(feat_dim, 128), nn.GELU(), nn.Linear(128, 2), nn.Sigmoid()
        )
    def forward(self, feats):
        return self.cls(feats), self.pos(feats)


@torch.no_grad()
def extract(model, x, dtype):
    raw = model.module if hasattr(model, 'module') else model
    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 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=30)
    ap.add_argument('--lr', type=float, default=5e-3)
    ap.add_argument('--pos-weight', type=float, default=0.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']
    n_cls = len(label_to_idx)
    img_size = get_arch_img_size(arch)
    bb = 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'}
    bb.load_state_dict(sd, strict=False)
    for p in bb.parameters(): p.requires_grad = False
    bb.eval()

    tr_ds = PosDS(args.manifest, label_to_idx, img_size, True, args.val_fold, args.min_samples)
    va_ds = PosDS(args.manifest, label_to_idx, img_size, False, args.val_fold, args.min_samples)
    log(f"Train={len(tr_ds)} Val={len(va_ds)}")

    cc = Counter(tr_ds.l2i[r['unified_label']] for r in tr_ds.records)
    sample_w = 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(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)

    with torch.no_grad():
        feat_dim = extract(bb, torch.zeros(1, 3, img_size, img_size, device=device), dtype).size(-1)
    head = MultiTaskHead(feat_dim, 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)

    best, best_probs, best_y = 0, None, None
    for ep in range(args.epochs):
        head.train()
        tl, nb = 0, 0
        for x, y, pos in tr_dl:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            pos = pos.to(device, non_blocking=True)
            f = extract(bb, x, dtype)
            logits_cls, pred_pos = head(f)
            loss_cls = F.cross_entropy(logits_cls, y, label_smoothing=0.1)
            loss_pos = F.mse_loss(pred_pos, pos)
            loss = loss_cls + args.pos_weight * loss_pos
            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)
                f = extract(bb, x, dtype)
                logits, _ = head(f)
                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"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},
                        args.output)
    log(f"=== MultiTask BEST: {best:.4f}")

if __name__ == '__main__':
    main()