#!/usr/bin/env python3
"""Smoke-test converted DeCUR checkpoints."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parent

MODELS = {
    "decur-resnet50-s1": {"channels": 2, "hidden": 2048, "seq": 49},
    "decur-resnet50-s2c": {"channels": 13, "hidden": 2048, "seq": 49},
    "decur-resnet50-rgb": {"channels": 3, "hidden": 2048, "seq": 49},
    "decur-resnet50-dem": {"channels": 3, "hidden": 2048, "seq": 49},
    "decur-resnet50-rda-s1": {"channels": 2, "hidden": 2048, "seq": 49},
    "decur-vit-small-patch16-s1": {"channels": 2, "hidden": 384, "seq": 197},
    "decur-vit-small-patch16-s2c": {"channels": 13, "hidden": 384, "seq": 197},
    "decur-vit-small-patch16-rgb": {"channels": 3, "hidden": 384, "seq": 197},
    "decur-mit-b2-rgb": {"channels": 3, "hidden": 512, "seq": 49},
    "decur-mit-b5-rgb": {"channels": 3, "hidden": 512, "seq": 49},
}


def load_model(model_dir: Path):
    from decur.models.decur.modeling_decur import DeCURModel

    return DeCURModel.from_pretrained(model_dir, local_files_only=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model",
        default="decur-mit-b2-rgb",
        choices=sorted(MODELS),
        help="Converted checkpoint folder name under DECUR-transformers/",
    )
    parser.add_argument("--all", action="store_true", help="Run all smoke tests")
    args = parser.parse_args()

    names = sorted(MODELS) if args.all else [args.model]
    for name in names:
        spec = MODELS[name]
        model_dir = ROOT / name
        model = load_model(model_dir)
        model.eval()

        x = torch.randn(1, spec["channels"], 224, 224)
        with torch.no_grad():
            out = model(pixel_values=x)

        pooled = tuple(out.pooler_output.shape)
        seq = tuple(out.last_hidden_state.shape)
        assert pooled == (1, spec["hidden"]), pooled
        assert seq == (1, spec["seq"], spec["hidden"]), seq
        print(f"OK {name}: pooler={pooled}, sequence={seq}")


if __name__ == "__main__":
    main()