"""Native CPU verification for the bounded Lance Wan2.2 VAE release."""
from __future__ import annotations

import hashlib
import json
import sys
import time
from pathlib import Path

import torch


ROOT = Path("outputs/unified_components_2026_10_20/lance_wan22_vae")


def sha256(path: Path) -> str:
    with path.open("rb") as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def main() -> None:
    started = time.monotonic()
    sys.path.insert(0, str(ROOT))
    from component import load_model

    source = torch.load(ROOT / "model.pth", map_location="cpu", weights_only=True)
    vae = load_model(ROOT, device="cpu", dtype=torch.float32)
    loaded = vae.model.state_dict()
    assert set(source) == set(loaded)
    assert all(torch.equal(source[name], loaded[name]) for name in source)
    source_hash = json.loads((ROOT / "source_tensor_sha256.json").read_text())["sha256"]
    checks = [
        {"name": "native_strict_load", "passed": True, "tensors": len(source)},
        {"name": "source_file_sha256", "passed": sha256(ROOT / "model.pth") == source_hash},
    ]
    # This 16x16 one-frame tensor is the smallest spatially valid 2x-patchified video.
    # The test validates VAE interfaces and geometry, not perceptual quality.
    x = torch.zeros((1, 3, 1, 16, 16), dtype=torch.float32)
    with torch.inference_mode():
        latent, log_var = vae.encode(x)
        reconstruction = vae.decode(latent)
    assert latent.shape[:2] == (1, 48)
    assert log_var.shape == latent.shape
    assert reconstruction.shape == x.shape
    assert torch.isfinite(latent).all() and torch.isfinite(reconstruction).all()
    checks.extend(
        [
            {"name": "encode_minimal_video", "passed": True, "input": list(x.shape), "latent": list(latent.shape)},
            {"name": "decode_minimal_video", "passed": True, "reconstruction": list(reconstruction.shape)},
            {"name": "finite_native_outputs", "passed": True},
        ]
    )
    report = {
        "strict_load": True,
        "source_tensor_bytes_equal": True,
        "checks": checks,
        "smoke": {"passed": True, "device": "cpu", "dtype": "float32"},
        "elapsed_seconds": time.monotonic() - started,
        "limitations": [
            "Native component boundary test only; the Lance MoT backbone and diffusion flow modules are excluded.",
            "The minimal CPU tensor verifies shape and finite outputs, not visual quality or long-video temporal quality.",
        ],
    }
    (ROOT / "validation.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()