"""Validate projection-strided Sage2 NHD execution on real H3 blocks."""

from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

import torch

from h3_blackwell_runtime.adaln import H3CurveAdaLN
from h3_blackwell_runtime.block import H3DiTBlock
from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.rope import h3_rope_rotation
from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas
from h3_blackwell_runtime.t2v import random_av_latents
from profile_h3_block import summarize


def timed(fn, warmup: int, iterations: int) -> dict[str, float]:
    with torch.inference_mode():
        for _ in range(warmup):
            fn()
        values = []
        for _ in range(iterations):
            torch.cuda.synchronize()
            started = time.perf_counter()
            fn()
            torch.cuda.synchronize()
            values.append(time.perf_counter() - started)
    return summarize(values)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--width", type=int, default=1344)
    parser.add_argument("--height", type=int, default=768)
    parser.add_argument("--frames", type=int, default=124)
    parser.add_argument("--steps", type=int, default=12)
    parser.add_argument("--sampler-step", type=int, default=1)
    parser.add_argument("--seed", type=int, default=440420)
    parser.add_argument("--text-tokens", type=int, default=100)
    parser.add_argument("--blocks", nargs="+", type=int, default=(0, 24, 49))
    parser.add_argument("--warmup", type=int, default=3)
    parser.add_argument("--iterations", type=int, default=10)
    parser.add_argument("--device", default="cuda")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)
    checkpoint = H3Checkpoint(args.model_path, device=args.device)
    packer = H3PromptPacker(checkpoint)
    video, audio, aligned_frames = random_av_latents(
        args.width, args.height, args.frames, args.seed, device=args.device,
    )
    sigma = beta_sigmas(args.steps, device=args.device)[args.sampler_step - 1]
    native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma)
    text = torch.randn(1, args.text_tokens, 5376, device=args.device, dtype=torch.bfloat16)
    hidden, timesteps, segments, positions, _, _ = packer(
        text, video, native_audio, _model_sigma(sigma),
    )
    rotation = h3_rope_rotation(
        positions.to(args.device), checkpoint.tensor("rope.inv_freq", dtype=torch.float32), hidden.dtype,
    )

    results = []
    original_layout = os.environ.get("H3_SAGE_QKV_LAYOUT")
    try:
        for block_index in args.blocks:
            if not 0 <= block_index < 50:
                raise ValueError("block indices must be in [0, 49]")
            block = H3DiTBlock.from_checkpoint(checkpoint, block_index, attention_backend="sage2").eval()
            block.fused_elementwise = True
            adaln = H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{block_index}.adaln_proj").eval()
            modulation = tuple(value.detach() for value in adaln(timesteps))
            run = lambda: block(hidden.clone(), rotation, *modulation, segments)

            with torch.inference_mode():
                os.environ["H3_SAGE_QKV_LAYOUT"] = "hnd"
                expected = run()
                os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd"
                actual = run()
            delta = actual.float() - expected.float()
            parity = {
                "equal": torch.equal(actual, expected),
                "max_abs": delta.abs().max().item(),
                "mean_abs": delta.abs().mean().item(),
                "reference_checksum": expected.float().sum().item(),
                "candidate_checksum": actual.float().sum().item(),
            }
            if not parity["equal"]:
                raise RuntimeError(f"block {block_index} Sage NHD output is not bit-exact: {parity}")

            os.environ["H3_SAGE_QKV_LAYOUT"] = "hnd"
            hnd_timing = timed(run, args.warmup, args.iterations)
            os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd"
            nhd_timing = timed(run, args.warmup, args.iterations)
            results.append({
                "block": block_index,
                "parity": parity,
                "hnd_timing": hnd_timing,
                "strided_nhd_timing": nhd_timing,
                "p50_speedup": hnd_timing["p50_s"] / nhd_timing["p50_s"],
                "p50_latency_reduction": 1.0 - nhd_timing["p50_s"] / hnd_timing["p50_s"],
            })
            print(
                f"block {block_index}: exact, hnd={hnd_timing['p50_s'] * 1000:.3f}ms, "
                f"strided_nhd={nhd_timing['p50_s'] * 1000:.3f}ms, speedup={results[-1]['p50_speedup']:.3f}x",
                flush=True,
            )
    finally:
        if original_layout is None:
            os.environ.pop("H3_SAGE_QKV_LAYOUT", None)
        else:
            os.environ["H3_SAGE_QKV_LAYOUT"] = original_layout

    report = {
        "device": torch.cuda.get_device_name(),
        "torch": torch.__version__,
        "resolution": [args.width, args.height],
        "frames": aligned_frames,
        "packed_tokens": hidden.shape[0],
        "steps": args.steps,
        "sampler_step": args.sampler_step,
        "seed": args.seed,
        "attention": "sage2",
        "fused_elementwise": True,
        "warmup": args.warmup,
        "iterations": args.iterations,
        "results": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()