# examples/vmamba_export.py
"""Export genuine VMamba-T predictions with the pinned official CPU reference.

Run from the repository after installing requirements-vmamba.txt in a venv:
  python examples/vmamba_export.py --cache-dir /tmp/vmamba-reference

The unmodified official model is imported from a hash-verified download. Only
its documented backend arguments are selected: torch scan, torch cross-scan.
No authored classifier, CUDA execution, or training is used by this exporter.
"""

import argparse
import ast
from datetime import datetime, timezone
from functools import partial
import hashlib
import importlib.metadata
import importlib.util
import json
from pathlib import Path
import platform
import sys
from types import SimpleNamespace
import urllib.request

from PIL import Image
import torch
from torchvision import transforms
from torchvision.models import ResNet50_Weights
from torchvision.transforms import functional as TF

REVISION = "2ed52ead062a51a64521ed3871d52914bf532876"
BASE_URL = f"https://raw.githubusercontent.com/MzeroMiko/VMamba/{REVISION}/"
REFERENCE_SHA256 = (
    "572207a2b6553ff85d81b602eebda2041623073f01fc0c6d1280c5a2add040f6"
)
CHECKPOINT_URL = (
    "https://github.com/MzeroMiko/VMamba/releases/download/"
    "%23v2cls/vssm1_tiny_0230s_ckpt_epoch_264.pth"
)
CHECKPOINT_SHA256 = (
    "d855b57224a652bc692804932f88dc3a39b9028e64be99f745a7fd6792ca672a"
)
MODEL_ID = "vmamba-t-s1l8"
TRANSFORM = {
    "resizeSize": 256,
    "cropSize": 224,
    "interpolation": "bicubic",
    "antialias": True,
    "mean": [0.485, 0.456, 0.406],
    "std": [0.229, 0.224, 0.225],
    "resizeRule": "shorter edge; preserve aspect ratio",
    "cropRule": "torchvision CenterCrop; round((dimension - 224) / 2)",
    "colorMode": "RGB",
    "tensorScale": "uint8 / 255",
    "source": BASE_URL + "classification/data/build.py",
    "settingsSource": BASE_URL + "classification/config.py",
}


def digest(path):
    """Return SHA-256 of exact file bytes without loading a checkpoint twice."""
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def verify_file(path, expected):
    """Reject changed cached input rather than treating it as pinned evidence."""
    if digest(path) != expected:
        raise ValueError(f"SHA-256 mismatch: {path}")


def cached_download(url, path, expected=None):
    """Download atomically, then verify known immutable inputs."""
    path = Path(path)
    if not path.exists():
        temporary = path.with_suffix(path.suffix + ".download")
        with urllib.request.urlopen(
            url, timeout=30
        ) as response, temporary.open("wb") as output:
            while chunk := response.read(1024 * 1024):
                output.write(chunk)
        temporary.replace(path)
    if expected:
        verify_file(path, expected)
    return path


def import_file(path, name):
    """Import a verified local Python file; timm registration needs sys.modules."""
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def load_official(source_path, checkpoint_path):
    """Strictly load the exact official Tiny classifier and force CPU backends."""
    verify_file(source_path, REFERENCE_SHA256)
    verify_file(checkpoint_path, CHECKPOINT_SHA256)
    official = import_file(source_path, "vmamba_official_pinned")
    model = official.vmamba_tiny_s1l8(pretrained=False, channel_first=True)
    checkpoint = torch.load(
        checkpoint_path, map_location="cpu", weights_only=True
    )
    loaded = model.load_state_dict(checkpoint["model"], strict=True)
    for block in model.modules():
        if isinstance(block, official.SS2D):
            block.forward_core = partial(
                block.forward_corev2,
                force_fp32=False,
                selective_scan_backend="torch",
                scan_force_torch=True,
            )
    model.eval()
    return (
        official,
        model,
        {
            "missingKeys": list(loaded.missing_keys),
            "unexpectedKeys": list(loaded.unexpected_keys),
            "strict": True,
        },
    )


def preprocess(image):
    """Return exact input pixels, normalized NCHW tensor and crop geometry."""
    resized = TF.resize(
        image.convert("RGB"),
        TRANSFORM["resizeSize"],
        interpolation=transforms.InterpolationMode.BICUBIC,
        antialias=True,
    )
    crop = TF.center_crop(resized, [224, 224])
    left = round((resized.width - 224) / 2)
    top = round((resized.height - 224) / 2)
    tensor = TF.normalize(
        TF.to_tensor(crop), TRANSFORM["mean"], TRANSFORM["std"]
    )
    return (
        crop,
        tensor.unsqueeze(0),
        {
            "resizedShape": [resized.height, resized.width],
            "cropBox": [left, top, left + 224, top + 224],
        },
    )


def official_transform(source_path):
    """Execute the pinned evaluation function independently of preprocess().

    The extracted function is unchanged. Its globals supply the actual torchvision
    transforms and the pinned config's default evaluation values; no training
    branches are called. The full original file stays available in the cache.
    """
    tree = ast.parse(Path(source_path).read_text())
    function = next(
        n
        for n in tree.body
        if isinstance(n, ast.FunctionDef) and n.name == "build_transform"
    )
    namespace = {
        "transforms": transforms,
        "_pil_interp": lambda name: transforms.InterpolationMode(name),
        "IMAGENET_DEFAULT_MEAN": (0.485, 0.456, 0.406),
        "IMAGENET_DEFAULT_STD": (0.229, 0.224, 0.225),
    }
    exec(
        compile(
            ast.Module(body=[function], type_ignores=[]),
            str(source_path),
            "exec",
        ),
        namespace,
    )
    config = SimpleNamespace(
        DATA=SimpleNamespace(IMG_SIZE=224, INTERPOLATION="bicubic"),
        TEST=SimpleNamespace(CROP=True),
    )
    return namespace["build_transform"](False, config)


def summarize_logits(logits, categories):
    """Keep full precision vectors and deterministic, class-index-tied top five."""
    if (
        logits.shape != (1000,)
        or len(categories) != 1000
        or not torch.isfinite(logits).all()
    ):
        raise ValueError("Expected exactly 1000 finite logits and class names")
    probabilities = logits.softmax(dim=0)
    indices = sorted(range(1000), key=lambda i: (-float(logits[i]), i))[:5]
    return {
        "logits": logits.tolist(),
        "probabilities": probabilities.tolist(),
        "top5": [
            {
                "index": i,
                "label": categories[i],
                "probability": float(probabilities[i]),
                "logit": float(logits[i]),
            }
            for i in indices
        ],
    }


def compare_teaching_scan(official, teaching_path):
    """Independently compare outputs AND gradients with unchanged official code."""
    teaching = import_file(teaching_path, "vmamba_teaching_export_check")
    torch.manual_seed(7)
    args = [
        torch.randn(2, 8, 9),
        torch.randn(2, 8, 9),
        -torch.rand(8, 3),
        torch.randn(2, 4, 3, 9),
        torch.randn(2, 4, 3, 9),
        torch.randn(8),
        torch.randn(8),
    ]
    ours = [value.clone().requires_grad_() for value in args]
    theirs = [value.clone().requires_grad_() for value in args]
    actual = teaching.selective_scan(*ours)
    expected = official.selective_scan_torch(*theirs)
    torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5)
    actual.square().sum().backward()
    expected.square().sum().backward()
    errors = []
    for left, right in zip(ours, theirs):
        torch.testing.assert_close(left.grad, right.grad, atol=2e-4, rtol=2e-4)
        errors.append(float((left.grad - right.grad).abs().max()))
    return {
        "reference": "unchanged official selective_scan_torch",
        "seed": 7,
        "inputShape": [2, 8, 9],
        "groups": 4,
        "stateSize": 3,
        "maxOutputAbsoluteError": float(
            (actual - expected).abs().max().detach()
        ),
        "maxGradientAbsoluteError": max(errors),
        "outputAtol": 2e-5,
        "outputRtol": 2e-5,
        "gradientAtol": 2e-4,
        "gradientRtol": 2e-4,
        "passed": True,
        "classifier": "official model executed directly; no authored port",
    }


def record_forward(model, tensor):
    """Capture actual module input/output shapes before each downsample."""
    stages = []
    handles = []

    def capture(name, repeats=None):
        def hook(_module, inputs, output):
            entry = {
                "name": name,
                "inputShape": list(inputs[0].shape),
                "outputShape": list(output.shape),
            }
            if repeats is not None:
                entry["repeats"] = repeats
            stages.append(entry)

        return hook

    modules = [("stem", model.patch_embed, None)]
    for index, layer in enumerate(model.layers):
        modules.append((f"stage{index + 1}", layer.blocks, len(layer.blocks)))
        if index < 3:
            modules.append((f"downsample{index + 1}", layer.downsample, None))
    modules += [
        ("pooled", model.classifier.flatten, None),
        ("logits", model.classifier.head, None),
    ]
    try:
        for name, module, repeats in modules:
            handles.append(module.register_forward_hook(capture(name, repeats)))
        with torch.inference_mode():
            logits = model(tensor)[0]
    finally:
        for handle in handles:
            handle.remove()
    return logits, stages


def export(cache, root, threads=2):
    """Publish a complete manifest only after all three real executions pass."""
    cache.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(threads)
    torch.manual_seed(0)
    torch.use_deterministic_algorithms(True)
    source = cached_download(
        BASE_URL + "vmamba.py", cache / "upstream-vmamba.py", REFERENCE_SHA256
    )
    checkpoint = cached_download(
        CHECKPOINT_URL, cache / "vmamba-t.pth", CHECKPOINT_SHA256
    )
    artifacts = [
        (
            "upstream-data-build.py",
            "classification/data/build.py",
            "c0865e9780bb1dd956f4a175239c6b53021bfd8ac264ecc5604dc07f6accdcb6",
        ),
        (
            "upstream-config.py",
            "classification/config.py",
            "1b6ba6868a3bde01308366c4f689c49c60dc3ef4b734ffcd100a950b8b389690",
        ),
        (
            "upstream-tiny.yaml",
            "classification/configs/vssm/vmambav2v_tiny_224.yaml",
            "6a7a22e91e4d74039d07a7b4cdf39da1b3510245d834ccebf74a003ffcb6d81b",
        ),
        (
            "upstream-LICENSE",
            "LICENSE",
            "fb8d4c2d5c5bc007e3176259bde6647bbb2a7a017eaab52f4a86e44d86c61076",
        ),
    ]
    source_artifacts = []
    for local, remote, expected in artifacts:
        path = cached_download(BASE_URL + remote, cache / local, expected)
        source_artifacts.append(
            {"path": remote, "url": BASE_URL + remote, "sha256": digest(path)}
        )
    official, model, strict_load = load_official(source, checkpoint)
    oracle = compare_teaching_scan(official, root / "examples/vmamba.py")
    evaluation_transform = official_transform(cache / "upstream-data-build.py")
    # torchvision's standard ImageNet-1K index vocabulary; no ResNet executes.
    categories = ResNet50_Weights.IMAGENET1K_V2.meta["categories"]
    source_photos = json.loads(
        (root / "examples/efficientnet-sample-images/manifest.json").read_text()
    )["photos"]
    output_dir = root / "public/images/papers/vmamba/examples"
    output_dir.mkdir(parents=True, exist_ok=True)
    runs = []
    for photo in source_photos:
        original = root / "public" / photo["originalImage"].lstrip("/")
        verify_file(original, photo["source"]["sha256"])
        with Image.open(original) as image:
            crop, tensor, geometry = preprocess(image)
            torch.testing.assert_close(
                tensor[0],
                evaluation_transform(image.convert("RGB")),
                atol=0,
                rtol=0,
            )
        temporary = output_dir / f".{photo['id']}-crop.png"
        crop.save(temporary, format="PNG")
        crop_hash = digest(temporary)
        crop_path = output_dir / f"{photo['id']}-{crop_hash[:12]}.png"
        temporary.replace(crop_path)
        print(
            f"Running official VMamba-T on {photo['id']} (CPU float32)",
            flush=True,
        )
        logits, stages = record_forward(model, tensor)
        run = {
            "id": photo["id"] + "-" + MODEL_ID,
            "photoId": photo["id"],
            "modelId": MODEL_ID,
            "inputImage": "/" + str(crop_path.relative_to(root / "public")),
            "inputSha256": crop_hash,
            "inputShape": list(tensor.shape),
            "normalizedTensorSha256": hashlib.sha256(
                tensor.numpy().tobytes()
            ).hexdigest(),
            "preprocessingReferenceMaxAbsoluteError": 0.0,
            **geometry,
            **summarize_logits(logits, categories),
            "stages": stages,
        }
        runs.append(run)
        print(
            json.dumps({"photo": photo["id"], "top5": run["top5"]}), flush=True
        )
    manifest = {
        "schemaVersion": 1,
        "provenance": {
            "sourceRevision": REVISION,
            "sourceUrl": BASE_URL + "vmamba.py",
            "referenceSha256": digest(source),
            "exportSha256": digest(Path(__file__)),
            "teachingSha256": digest(root / "examples/vmamba.py"),
            "exportedAt": datetime.now(timezone.utc).isoformat(),
            "device": "cpu",
            "dtype": "float32",
            "numThreads": threads,
            "deterministicAlgorithms": True,
            "backend": "official PyTorch reference",
            "versions": {
                "python": platform.python_version(),
                **{
                    name: importlib.metadata.version(name)
                    for name in [
                        "torch",
                        "torchvision",
                        "Pillow",
                        "timm",
                        "fvcore",
                        "numpy",
                    ]
                },
            },
            "platform": platform.platform(),
            "oracleAgreement": oracle,
            "preprocessingReference": "unchanged pinned build_transform(False), tensor equality on every photo",
            "sourceArtifacts": source_artifacts,
            "codeLicense": {"name": "MIT", "url": BASE_URL + "LICENSE"},
            "labelsSource": "torchvision.models.ResNet50_Weights.IMAGENET1K_V2.meta.categories (vocabulary only)",
        },
        "categories": categories,
        "models": [
            {
                "id": MODEL_ID,
                "label": "VMamba-T (s1l8)",
                "factory": "vmamba_tiny_s1l8",
                "strictLoad": strict_load,
                "parameterCount": sum(
                    parameter.numel() for parameter in model.parameters()
                ),
                "weights": {
                    "url": CHECKPOINT_URL,
                    "sha256": digest(checkpoint),
                    "bytes": checkpoint.stat().st_size,
                    "key": "model",
                },
                "config": {
                    "depths": [2, 2, 8, 2],
                    "dims": [96, 192, 384, 768],
                    "ssmDState": 1,
                    "ssmRatio": 1,
                    "ssmDtRank": "auto",
                    "ssmConv": 3,
                    "ssmConvBias": False,
                    "forwardType": "v05_noz",
                    "mlpRatio": 4,
                    "normLayer": "ln2d",
                    "patchEmbed": "v2",
                    "downsample": "v3",
                    "numClasses": 1000,
                },
                "transform": TRANSFORM,
            }
        ],
        "photos": source_photos,
        "runs": runs,
        "limitations": [
            "Three individual photo predictions, not an accuracy evaluation.",
            "No latency measurement, attention map, or saliency claim.",
            "CPU float32 reference results; released GPU kernels may differ numerically.",
        ],
    }
    target = root / "public/data/vmamba/predictions.json"
    target.parent.mkdir(parents=True, exist_ok=True)
    temporary = target.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(manifest, indent=2, allow_nan=False) + "\n")
    temporary.replace(target)
    print(f"Wrote {target}", flush=True)


def main():
    """Parse only reproducibility controls; architecture and weights stay pinned."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-dir", type=Path, required=True)
    parser.add_argument(
        "--repo-root", type=Path, default=Path(__file__).resolve().parents[1]
    )
    parser.add_argument("--threads", type=int, default=2)
    args = parser.parse_args()
    if args.threads < 1:
        parser.error("--threads must be positive")
    export(args.cache_dir, args.repo_root, args.threads)


if __name__ == "__main__":
    main()