#!/usr/bin/env python3
"""
predict_nexar_test.py
═══════════════════════════════════════════════════════════════════════════════
Generate a Kaggle-style submission CSV from a Nexar head checkpoint and the
test belief cache.

Inputs:
  --head_ckpt   checkpoints/Nexar/qwen3vl4b_head/best.pt
  --test_cache  data/belief_cache_nexar_qwen3vl4b/test.pt

Output CSV columns:  id,score

Usage
─────
  python -m training.Policy.predict_nexar_test \
      --head_ckpt  checkpoints/Nexar/qwen3vl4b_head/best.pt \
      --test_cache data/belief_cache_nexar_qwen3vl4b/test.pt \
      --out        submissions/nexar_qwen3vl4b.csv
"""
from __future__ import annotations

import argparse
import csv
import logging
from pathlib import Path

import numpy as np
import torch

from training.Policy.train_nexar_head import NexarHead

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.predict_nexar_test")


def main():
    ap = argparse.ArgumentParser("predict_nexar_test")
    ap.add_argument("--head_ckpt",  required=True)
    ap.add_argument("--test_cache", required=True)
    ap.add_argument("--out",        required=True)
    ap.add_argument("--batch_size", type=int, default=128)
    args = ap.parse_args()

    logger.info(f"loading head {args.head_ckpt}")
    ck = torch.load(args.head_ckpt, map_location="cpu", weights_only=False)
    meta = ck["meta"]
    model = NexarHead(hidden_dim=meta["hidden_dim"],
                      proj_dim=meta["proj_dim"],
                      n_layers=meta["n_layers"],
                      n_heads=meta["n_heads"],
                      dropout=meta["dropout"])
    model.load_state_dict(ck["state_dict"])
    model.eval().to("cuda")

    logger.info(f"loading test cache {args.test_cache}")
    te = torch.load(args.test_cache, map_location="cpu", weights_only=False)
    x = te["beliefs_frame"].float()
    v = te["valid_frames"].bool()
    ids = te["meta"]["video_ids"]
    assert x.shape[0] == len(ids), f"cache/ids mismatch: {x.shape[0]} vs {len(ids)}"

    probs = []
    with torch.no_grad():
        for i in range(0, x.size(0), args.batch_size):
            xb = x[i:i + args.batch_size].to("cuda")
            vb = v[i:i + args.batch_size].to("cuda")
            logits = model(xb, vb).cpu().numpy()
            probs.append(1 / (1 + np.exp(-logits)))
    probs = np.concatenate(probs)
    assert len(probs) == len(ids)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["id", "score"])
        for vid, p in zip(ids, probs):
            w.writerow([vid, f"{float(p):.6f}"])
    logger.info(f"wrote {len(ids)} rows -> {out}")


if __name__ == "__main__":
    main()