#!/usr/bin/env python3
"""Smoke-test converted Prithvi-EO-2.0 checkpoints."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch

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

MODELS = {
    "prithvi-eo-v2-tiny-tl": {"hidden": 192, "seq": 785},
    "prithvi-eo-v2-100m-tl": {"hidden": 768, "seq": 785},
    "prithvi-eo-v2-300m": {"hidden": 1024, "seq": 785},
    "prithvi-eo-v2-300m-tl": {"hidden": 1024, "seq": 785},
    "prithvi-eo-v2-600m": {"hidden": 1280, "seq": 1025},
    "prithvi-eo-v2-600m-tl": {"hidden": 1280, "seq": 1025},
}


def load_model(model_dir: Path):
    from transformers import AutoModel

    return AutoModel.from_pretrained(model_dir, trust_remote_code=True, local_files_only=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="prithvi-eo-v2-300m-tl", choices=sorted(MODELS))
    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()

        pixel_values = torch.randn(1, 6, 4, 224, 224)
        temporal_coords = torch.tensor([[[2018, 26], [2018, 106], [2018, 201], [2018, 266]]], dtype=torch.float32)
        location_coords = torch.tensor([[19.5, -99.1]], dtype=torch.float32)

        with torch.no_grad():
            out = model(
                pixel_values=pixel_values,
                temporal_coords=temporal_coords if "tl" in name else None,
                location_coords=location_coords if "tl" in name else None,
            )

        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

        frames = [np.random.uniform(500, 3000, (224, 224, 6)).astype("float32") for _ in range(4)]
        from transformers import pipeline

        pipe = pipeline(
            task="prithvi-eo-feature-extraction",
            model=model_dir,
            trust_remote_code=True,
            local_files_only=True,
        )
        features = pipe(
            frames,
            pool=True,
            return_tensors=True,
            temporal_coords=[[2018, 26], [2018, 106], [2018, 201], [2018, 266]] if "tl" in name else None,
            location_coords=[19.5, -99.1] if "tl" in name else None,
        )
        assert tuple(features.shape) == (1, spec["hidden"]), tuple(features.shape)
        print(f"OK {name}: pooler={pooled}, sequence={seq}, pipeline={tuple(features.shape)}")


if __name__ == "__main__":
    main()