"""Compare baseline and fused-modulation H3 sampling in one resident model."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.lora import load_lora_adapter, set_active_lora
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.sampler import sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents
from h3_blackwell_runtime.token_refiner import H3TokenRefiner


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=2)
    parser.add_argument("--warmup-steps", type=int, default=1)
    parser.add_argument("--seed", type=int, default=440420)
    parser.add_argument("--text-tokens", type=int, default=100)
    parser.add_argument("--attention", default="sage2")
    parser.add_argument(
        "--feature", choices=("modulate_fusion", "swiglu_fusion", "lora_producer_fusion"),
        default="modulate_fusion",
    )
    parser.add_argument("--lora-path")
    parser.add_argument("--lora-name", default="validation")
    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)
    model = H3PackedDenoiser.from_checkpoint(
        checkpoint, output_dtype=torch.bfloat16, attention_backend=args.attention,
    ).eval()
    refiner = None
    if args.lora_path:
        refiner = H3TokenRefiner(checkpoint, attention_backend=args.attention).eval()
        load_lora_adapter(model, refiner, args.lora_name, args.lora_path, args.device)
        set_active_lora(model, refiner, args.lora_name)
    packer = H3PromptPacker(checkpoint)
    if hasattr(checkpoint, "release_cache"):
        checkpoint.release_cache()
    video, audio, aligned_frames = random_av_latents(
        args.width, args.height, args.frames, args.seed, device=args.device,
    )
    text = torch.randn(
        1, args.text_tokens, 5376, device=args.device, dtype=torch.bfloat16,
    )

    def run(enabled: bool, steps: int):
        for block in model.backbone.blocks:
            if args.feature == "lora_producer_fusion":
                block.fused_nvfp4_modulation = enabled
                block.mlp.fused_nvfp4_swiglu = enabled
            elif args.feature == "swiglu_fusion":
                block.mlp.fused_nvfp4_swiglu = enabled
            else:
                block.fused_nvfp4_modulation = enabled
        torch.cuda.synchronize()
        started = time.perf_counter()
        result = sample_video_res_multistep(
            model,
            packer,
            text,
            video.clone(),
            audio.clone(),
            steps=steps,
            seed=args.seed,
            return_audio=True,
            progress=True,
        )
        torch.cuda.synchronize()
        return result, time.perf_counter() - started

    with torch.inference_mode():
        if args.warmup_steps:
            run(False, args.warmup_steps)
            run(True, args.warmup_steps)
        (reference_video, reference_audio), baseline_s = run(False, args.steps)
        (candidate_video, candidate_audio), candidate_s = run(True, args.steps)

    video_delta = candidate_video.float() - reference_video.float()
    audio_delta = candidate_audio.float() - reference_audio.float()
    report = {
        "device": torch.cuda.get_device_name(),
        "resolution": [args.width, args.height],
        "frames": aligned_frames,
        "steps": args.steps,
        "feature": args.feature,
        "seed": args.seed,
        "baseline_seconds": baseline_s,
        "candidate_seconds": candidate_s,
        "improvement_percent": (1.0 - candidate_s / baseline_s) * 100.0,
        "video_equal": torch.equal(candidate_video, reference_video),
        "audio_equal": torch.equal(candidate_audio, reference_audio),
        "video_max_abs": video_delta.abs().max().item(),
        "audio_max_abs": audio_delta.abs().max().item(),
        "reference_checksums": [
            reference_video.float().sum().item(), reference_audio.float().sum().item(),
        ],
        "candidate_checksums": [
            candidate_video.float().sum().item(), candidate_audio.float().sum().item(),
        ],
    }
    report["equal"] = report["video_equal"] and report["audio_equal"]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2), flush=True)
    if not report["equal"]:
        raise RuntimeError("fused modulation trajectory is not bit-exact")


if __name__ == "__main__":
    main()