#!/usr/bin/env python3
"""Benchmark L1 visual model latency on decoded RGB histories.

The benchmark excludes mp4 decoding: it first materializes a small official-test
RGB batch in CPU memory, then times model forward passes including CPU->GPU
transfer, normalization/resizing, backbone, temporal/head layers, and GPU
synchronization.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any, Dict, List

import torch
import torch.nn as nn
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "scripts"))

from build_rgb_frame_cache import RGBVideoIndex, parse_cell_frames  # noqa: E402
from eval_l1_foundation_baselines import FrozenBackbone, MultiHeadProbe  # noqa: E402
from eval_l1_latent_understanding import FIELDS, parse_offsets  # noqa: E402
from layered_belief import collect_decision_samples  # noqa: E402
from train_rgb_situation_net import (  # noqa: E402
    RGBSituationNet,
    encode_state,
    normalize_rgb,
    state_batch,
)


def load_manifest(path: Path) -> Dict[str, Any]:
    with path.open(encoding="utf-8") as f:
        return json.load(f)


def load_rgb_histories(args: argparse.Namespace, rows: List[Dict[str, Any]]) -> torch.Tensor:
    offsets = parse_offsets(args.history_offsets)
    frame_in_cell = parse_cell_frames(args.frame_in_cell)
    index = RGBVideoIndex(
        args.processed_root,
        offsets,
        frame_in_cell,
        args.height,
        args.width,
        args.backend,
        args.max_open,
    )
    xs: List[torch.Tensor] = []
    for row in rows:
        hist, _ = index.history(row["boss"], int(row["fight"]), row["belief"]["time"])
        if hist is None:
            raise RuntimeError(f"missing RGB history for {row['boss']} fight{row['fight']} t={row['belief']['time']}")
        xs.append(hist)
    index.close()
    return torch.stack(xs, dim=0).contiguous()


def sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def timed_loop(fn, n: int, batch_size: int, warmup: int, repeats: int, device: torch.device) -> Dict[str, float]:
    with torch.no_grad():
        for _ in range(warmup):
            for start in range(0, n, batch_size):
                fn(start, min(start + batch_size, n))
        sync(device)
        t0 = time.perf_counter()
        for _ in range(repeats):
            for start in range(0, n, batch_size):
                fn(start, min(start + batch_size, n))
        sync(device)
    elapsed = time.perf_counter() - t0
    decisions = n * repeats
    return {
        "seconds_total": round(elapsed, 6),
        "decisions": int(decisions),
        "ms_per_decision": round(1000.0 * elapsed / max(1, decisions), 4),
        "decisions_per_second": round(decisions / max(1e-9, elapsed), 2),
    }


def benchmark_task_model(
    checkpoint: str,
    batch_size: int,
    rgb: torch.Tensor,
    rows: List[Dict[str, Any]],
    args: argparse.Namespace,
    device: torch.device,
) -> Dict[str, Any]:
    ckpt = torch.load(checkpoint, map_location=device, weights_only=False)
    ckpt_args = ckpt.get("args", {})
    label_values = ckpt["label_values"]
    state_maps = ckpt["state_maps"]
    model = RGBSituationNet(
        {field: len(label_values[field]) for field in FIELDS},
        state_maps,
        backbone=ckpt_args.get("backbone", "resnet18"),
        pretrained=False,
        freeze_backbone=False,
        hidden=int(ckpt_args.get("hidden", 384)),
        state_mode=ckpt_args.get("state_mode", "boss"),
        spatial_tokens=int(ckpt_args.get("spatial_tokens", 0)),
        spatial_heads=int(ckpt_args.get("spatial_heads", 8)),
    ).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    state = encode_state(rows, state_maps)

    def forward(start: int, end: int) -> None:
        sl = torch.arange(start, end)
        xb = normalize_rgb(rgb[sl], device)
        sb = state_batch(state, sl, device)
        with torch.amp.autocast("cuda", enabled=args.amp and device.type == "cuda", dtype=torch.bfloat16):
            _ = model(xb, sb)

    out = timed_loop(forward, len(rows), batch_size, args.warmup, args.repeats, device)
    out.update({
        "batch_size": batch_size,
        "includes_video_decode": False,
        "checkpoint": str(Path(checkpoint)),
        "backbone": ckpt_args.get("backbone", "resnet18"),
        "training_mode": "end_to_end" if not ckpt_args.get("freeze_backbone", False) else "frozen_backbone",
    })
    return out


def foundation_features_on_gpu(backbone: FrozenBackbone, rgb: torch.Tensor, device: torch.device) -> torch.Tensor:
    if backbone.kind == "dinov2_small":
        b, t = rgb.shape[:2]
        frames = rgb.reshape(b * t, *rgb.shape[2:])
        x = frames.to(device, non_blocking=True).float() / 255.0
        x = F.interpolate(x, size=(backbone.image_size, backbone.image_size), mode="bilinear", align_corners=False)
        from train_rgb_situation_net import IMAGENET_MEAN, IMAGENET_STD

        mean = IMAGENET_MEAN.reshape(1, 3, 1, 1).to(device=device, dtype=x.dtype)
        std = IMAGENET_STD.reshape(1, 3, 1, 1).to(device=device, dtype=x.dtype)
        x = (x - mean) / std
        z = backbone.model(x).float().reshape(b, t, -1)
        return torch.cat([z.mean(dim=1), z[:, -1]], dim=-1)

    if backbone.kind == "videomae_base":
        b, t = rgb.shape[:2]
        frames = rgb.reshape(b * t, *rgb.shape[2:])
        x = frames.to(device, non_blocking=True).float() / 255.0
        x = F.interpolate(x, size=(backbone.image_size, backbone.image_size), mode="bilinear", align_corners=False)
        from train_rgb_situation_net import IMAGENET_MEAN, IMAGENET_STD

        mean = IMAGENET_MEAN.reshape(1, 3, 1, 1).to(device=device, dtype=x.dtype)
        std = IMAGENET_STD.reshape(1, 3, 1, 1).to(device=device, dtype=x.dtype)
        x = (x - mean) / std
        x = x.reshape(b, t, 3, backbone.image_size, backbone.image_size)
        idx = torch.linspace(0, t - 1, steps=16, device=x.device).round().long()
        x = x[:, idx]
        return backbone.model(pixel_values=x).last_hidden_state.float().mean(dim=1)

    raise AssertionError(backbone.kind)


def benchmark_foundation(
    name: str,
    rgb: torch.Tensor,
    rows: List[Dict[str, Any]],
    args: argparse.Namespace,
    device: torch.device,
    batch_size: int,
) -> Dict[str, Any]:
    backbone = FrozenBackbone(name, device, args.image_size)
    label_sizes = {
        "player_distance_bin": 3,
        "player_angle_bin": 3,
        "prev_boss_skill": 12,
        "hp_phase": 3,
    }
    bosses = {boss: i for i, boss in enumerate(sorted({row["boss"] for row in rows}))}
    boss = torch.tensor([bosses[row["boss"]] for row in rows], dtype=torch.long)
    probe = MultiHeadProbe(backbone.out_dim, len(bosses), label_sizes, args.hidden).to(device).eval()

    def forward(start: int, end: int) -> None:
        feat = foundation_features_on_gpu(backbone, rgb[start:end], device)
        _ = probe(feat, boss[start:end].to(device, non_blocking=True))

    out = timed_loop(forward, len(rows), batch_size, args.warmup, args.repeats, device)
    out.update({"batch_size": batch_size, "includes_video_decode": False})
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default=str(ROOT / "out/layered/manifest.json"))
    ap.add_argument("--processed_root", default=str(ROOT / "data/processed"))
    ap.add_argument("--pact_ckpt", default=str(ROOT / "out/situation/raw_rgb_resnet18_decision_finetune_from_dp_current0_e10.pt"))
    ap.add_argument("--e2e_dinov2_ckpt", default="")
    ap.add_argument("--out", default=str(ROOT / "out/situation/l1_visual_latency_benchmark.json"))
    ap.add_argument("--split", default="test")
    ap.add_argument("--n_samples", type=int, default=256)
    ap.add_argument("--history_offsets", default="-8,-4,-2,-1,0")
    ap.add_argument("--frame_in_cell", default="2")
    ap.add_argument("--height", type=int, default=192)
    ap.add_argument("--width", type=int, default=336)
    ap.add_argument("--image_size", type=int, default=224)
    ap.add_argument("--backend", choices=["auto", "decord", "opencv"], default="decord")
    ap.add_argument("--max_open", type=int, default=4)
    ap.add_argument("--pact_batch_size", type=int, default=16)
    ap.add_argument("--e2e_dinov2_batch_size", type=int, default=1)
    ap.add_argument("--dinov2_batch_size", type=int, default=16)
    ap.add_argument("--videomae_batch_size", type=int, default=8)
    ap.add_argument("--warmup", type=int, default=2)
    ap.add_argument("--repeats", type=int, default=5)
    ap.add_argument("--hidden", type=int, default=512)
    ap.add_argument("--amp", action="store_true")
    ap.add_argument(
        "--skip_foundation_baselines",
        action="store_true",
        help="Benchmark only task-trained checkpoints (PACT and optional end-to-end DINOv2).",
    )
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()

    device = torch.device(args.device)
    rows = collect_decision_samples(load_manifest(Path(args.manifest)), args.split, None)[: args.n_samples]
    t0 = time.perf_counter()
    rgb = load_rgb_histories(args, rows)
    decode_seconds = time.perf_counter() - t0
    methods = {
        "PACT-L1": benchmark_task_model(
            args.pact_ckpt, args.pact_batch_size, rgb, rows, args, device
        ),
    }
    if args.e2e_dinov2_ckpt:
        methods["DINOv2-S end-to-end"] = benchmark_task_model(
            args.e2e_dinov2_ckpt,
            args.e2e_dinov2_batch_size,
            rgb,
            rows,
            args,
            device,
        )
    if not args.skip_foundation_baselines:
        methods.update({
            "DINOv2-S + MLP": benchmark_foundation(
                "dinov2_small", rgb, rows, args, device, args.dinov2_batch_size
            ),
            "VideoMAE-B + MLP": benchmark_foundation(
                "videomae_base", rgb, rows, args, device, args.videomae_batch_size
            ),
        })
    result = {
        "definition": "model forward latency on pre-decoded RGB histories; excludes mp4 decoding",
        "device": str(device),
        "device_name": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
        "amp": bool(args.amp and device.type == "cuda"),
        "n_samples": len(rows),
        "warmup": args.warmup,
        "repeats": args.repeats,
        "history_offsets": parse_offsets(args.history_offsets),
        "frame_in_cell": parse_cell_frames(args.frame_in_cell),
        "rgb_shape": list(rgb.shape[1:]),
        "decode_seconds_excluded": round(decode_seconds, 4),
        "methods": methods,
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps(result, indent=2), flush=True)
    print(f"wrote {out}", flush=True)


if __name__ == "__main__":
    main()