#!/usr/bin/env python3
"""
Sanity tests for make_belief_cache_v2.py outputs.

Runs four checks per cache_mode (on already-built debug caches):
  T1 shape + dtype invariants
  T2 NaN/Inf-free
  T3 index alignment with manifest (n_samples == len(manifest['samples'][:N]))
  T4 cache_mode-specific semantic checks
       mean_pool   : must approximately match legacy v1 cache
                     (fp16 round-trip → atol=2e-2 cosine)
       dual_pool   : beliefs_img and beliefs_text differ; both unit-non-zero
       per_frame   : valid_frames & beliefs_frame consistent (zero rows ⟺ False)
       spatial4x4  : beliefs_grid[b,f] has 16 distinct rows w/ non-trivial variance
                     when valid_frames[b,f]=True
       (also: tta_means match v1 cache to ~1e-4 — TTA head is deterministic)

Usage
─────
  # Build debug caches first
  python -m training.Policy.make_belief_cache_v2 \\
      --sft_checkpoint checkpoints/SFT/sft_v2/best \\
      --cache_mode mean_pool   --debug --overwrite \\
      --out_dir data/belief_cache_v2_debug
  python -m training.Policy.make_belief_cache_v2 \\
      --sft_checkpoint checkpoints/SFT/sft_v2/best \\
      --cache_mode dual_pool   --debug --overwrite \\
      --out_dir data/belief_cache_v2_debug
  python -m training.Policy.make_belief_cache_v2 \\
      --sft_checkpoint checkpoints/SFT/sft_v2/best \\
      --cache_mode per_frame   --debug --overwrite \\
      --out_dir data/belief_cache_v2_debug
  python -m training.Policy.make_belief_cache_v2 \\
      --sft_checkpoint checkpoints/SFT/sft_v2/best \\
      --cache_mode spatial4x4  --debug --overwrite \\
      --out_dir data/belief_cache_v2_debug

  # Then run tests
  python -m training.Policy.test_belief_cache_v2 \\
      --cache_root data/belief_cache_v2_debug \\
      --legacy_cache data/belief_cache/val.pt \\
      --label_dir   data/policy_labels \\
      --n 16
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict

import torch
import torch.nn.functional as F


def _load(p: Path) -> Dict:
    return torch.load(p, map_location="cpu", weights_only=False)


def _check_no_nan_inf(name: str, t: torch.Tensor):
    if not t.dtype.is_floating_point:
        return
    nans = int(torch.isnan(t).sum().item())
    infs = int(torch.isinf(t).sum().item())
    assert nans == 0, f"{name} has {nans} NaN"
    assert infs == 0, f"{name} has {infs} Inf"


def _check_dtype(name: str, t: torch.Tensor, expect: torch.dtype):
    assert t.dtype == expect, f"{name} dtype={t.dtype}, expected {expect}"


def _ok(msg: str):
    print(f"  PASS  {msg}")


def _hdr(s: str):
    print(f"\n── {s} " + "─" * (78 - len(s)))


def test_mean_pool(cache_root: Path, legacy_path: Path, n_expected: int):
    _hdr("mean_pool")
    p = cache_root / "mean_pool" / "val.pt"
    if not p.exists():
        print(f"  SKIP  {p} not built")
        return
    d = _load(p)
    b = d["beliefs"]
    tm = d["tta_means"]
    tv = d["tta_vars"]

    # T1 shape + dtype
    assert b.dim() == 2,                f"beliefs dim={b.dim()}"
    assert b.shape[0] == n_expected,    f"beliefs N={b.shape[0]} vs {n_expected}"
    _check_dtype("beliefs",   b,  torch.float16)
    _check_dtype("tta_means", tm, torch.float32)
    _check_dtype("tta_vars",  tv, torch.float32)
    _ok(f"shapes/dtypes  beliefs={tuple(b.shape)} fp16, tta fp32")

    # T2 NaN/Inf
    for k in ("beliefs", "tta_means", "tta_vars"):
        _check_no_nan_inf(k, d[k])
    _ok("no NaN/Inf")

    # T4 vs legacy v1
    if legacy_path.exists():
        v1 = _load(legacy_path)
        v1_b   = v1["beliefs"][:n_expected].float()
        v1_tm  = v1["tta_means"][:n_expected].float()
        cos = F.cosine_similarity(b.float(), v1_b, dim=-1).mean().item()
        tta_diff = (tm.float() - v1_tm).abs().mean().item()
        assert cos > 0.95, f"v1↔v2 belief cosine={cos:.4f} (<0.95) — pooling logic differs"
        assert tta_diff < 1e-2, f"tta_mean diff vs v1: {tta_diff:.6f} (head should be deterministic)"
        _ok(f"matches legacy v1: belief_cos={cos:.4f}, tta_mae={tta_diff:.2e}")
    else:
        print(f"  SKIP  legacy comparison (no {legacy_path})")


def test_dual_pool(cache_root: Path, n_expected: int):
    _hdr("dual_pool")
    p = cache_root / "dual_pool" / "val.pt"
    if not p.exists():
        print(f"  SKIP  {p} not built")
        return
    d = _load(p)
    bi = d["beliefs_img"]
    bt = d["beliefs_text"]
    assert bi.shape == bt.shape and bi.shape[0] == n_expected
    _check_dtype("beliefs_img",  bi, torch.float16)
    _check_dtype("beliefs_text", bt, torch.float16)
    _ok(f"shapes/dtypes  img={tuple(bi.shape)} text={tuple(bt.shape)} fp16")

    _check_no_nan_inf("beliefs_img",  bi)
    _check_no_nan_inf("beliefs_text", bt)
    _ok("no NaN/Inf")

    # img and text means should differ — if identical, splitting failed
    cos_it = F.cosine_similarity(bi.float(), bt.float(), dim=-1).mean().item()
    assert cos_it < 0.999, f"img/text means too similar (cos={cos_it:.4f}) — split broken"
    norm_i = bi.float().norm(dim=-1).mean().item()
    norm_t = bt.float().norm(dim=-1).mean().item()
    assert norm_i > 1.0 and norm_t > 1.0, \
        f"degenerate norms img={norm_i:.3f} text={norm_t:.3f}"
    _ok(f"img/text differ: cos={cos_it:.4f}, norms img={norm_i:.2f} text={norm_t:.2f}")


def test_per_frame(cache_root: Path, n_expected: int):
    _hdr("per_frame")
    p = cache_root / "per_frame" / "val.pt"
    if not p.exists():
        print(f"  SKIP  {p} not built")
        return
    d = _load(p)
    bf = d["beliefs_frame"]    # [N, F, D]
    vf = d["valid_frames"]     # [N, F] bool
    bt = d["beliefs_text"]
    assert bf.dim() == 3 and bf.shape[0] == n_expected
    assert vf.shape == bf.shape[:2]
    assert vf.dtype == torch.bool
    _check_dtype("beliefs_frame", bf, torch.float16)
    _ok(f"shapes/dtypes  frame={tuple(bf.shape)} valid={tuple(vf.shape)}")

    _check_no_nan_inf("beliefs_frame", bf)
    _check_no_nan_inf("beliefs_text",  bt)
    _ok("no NaN/Inf")

    # invalid frames should be all-zero; valid frames non-zero
    invalid_norms = bf[~vf].float().norm(dim=-1)
    valid_norms   = bf[ vf].float().norm(dim=-1)
    if invalid_norms.numel() > 0:
        assert invalid_norms.max().item() < 1e-3, \
            f"invalid-frame slot has nonzero belief (max norm={invalid_norms.max():.3e})"
    if valid_norms.numel() > 0:
        assert valid_norms.min().item() > 0.5, \
            f"valid frame degenerate (min norm={valid_norms.min():.3e})"
    _ok(f"validity mask consistent: {int(vf.sum())} valid / {vf.numel()} slots")


def test_spatial4x4(cache_root: Path, n_expected: int):
    _hdr("spatial4x4")
    p = cache_root / "spatial4x4" / "val.pt"
    if not p.exists():
        print(f"  SKIP  {p} not built")
        return
    d = _load(p)
    bg = d["beliefs_grid"]    # [N, F, 16, D]
    vf = d["valid_frames"]    # [N, F]
    assert bg.dim() == 4 and bg.shape[0] == n_expected and bg.shape[2] == 16
    assert vf.shape == bg.shape[:2]
    _check_dtype("beliefs_grid", bg, torch.float16)
    _ok(f"shapes/dtypes  grid={tuple(bg.shape)} (B,F,16,D)")

    _check_no_nan_inf("beliefs_grid", bg)
    _ok("no NaN/Inf")

    # For valid frames, the 16 spatial cells should have non-trivial variance —
    # if all 16 are identical, the spatial pool collapsed somewhere.
    if vf.any():
        v_idx = vf.nonzero(as_tuple=False)   # [n_valid, 2]
        # sample up to 8 random (b,f) and compute variance across the 16 cells
        sel = v_idx[:min(8, v_idx.shape[0])]
        spatial_stds = []
        for (b, f) in sel.tolist():
            cells = bg[b, f].float()                      # [16, D]
            # std across the 16 cells, averaged over D
            spatial_stds.append(cells.std(dim=0).mean().item())
        avg_std = float(sum(spatial_stds) / max(len(spatial_stds), 1))
        assert avg_std > 1e-3, \
            f"spatial cells nearly constant (mean std across 16 cells={avg_std:.2e})"
        _ok(f"spatial variance OK (avg cross-cell std={avg_std:.3f})")

    # invalid frames must be zero
    inv_max = bg[~vf].float().abs().max().item() if (~vf).any() else 0.0
    assert inv_max < 1e-3, f"invalid frame slot nonzero (|max|={inv_max:.3e})"
    _ok("invalid frame slots are zero")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache_root", required=True, type=Path,
                    help="Root containing {mean_pool,dual_pool,per_frame,spatial4x4}/{train,val}.pt")
    ap.add_argument("--legacy_cache", default="data/belief_cache/val.pt", type=Path,
                    help="v1 cache for cross-check (used by mean_pool test)")
    ap.add_argument("--label_dir", default="data/policy_labels", type=Path)
    ap.add_argument("--n", type=int, default=16,
                    help="Expected n_samples (must match --debug_samples used at build time)")
    args = ap.parse_args()

    # Cross-verify expected N against label manifest
    val_labels = json.loads((args.label_dir / "val.json").read_text())
    n_total = len(val_labels.get("samples", []))
    n_expect = min(args.n, n_total)
    print(f"Expecting {n_expect} samples (debug_samples={args.n}, manifest has {n_total}).")

    test_mean_pool(args.cache_root, args.legacy_cache, n_expect)
    test_dual_pool(args.cache_root, n_expect)
    test_per_frame(args.cache_root, n_expect)
    test_spatial4x4(args.cache_root, n_expect)

    print("\nAll requested tests passed.")


if __name__ == "__main__":
    main()