#!/usr/bin/env python3
"""Smoke-test FireWx-FM checkpoint loading and tensor inference.

This test intentionally uses a synthetic 16-channel tensor. It verifies that
the released model code and checkpoint are mutually compatible; scientific
reproduction still requires the source datasets documented in data_sources/.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from firewxfm.modeling_unet import UNetSmallFlex
from firewxfm.serve_conus import load_checkpoint


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("models/checkpoints/firewxfm_2024_staticfix_region_balanced_bce_seed42.pt"),
    )
    parser.add_argument("--height", type=int, default=64)
    parser.add_argument("--width", type=int, default=64)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    device = torch.device(args.device)
    torch.manual_seed(0)
    model = UNetSmallFlex(
        in_ch=16,
        base=32,
        dropout=0.1,
        norm_type="group",
        norm_groups=8,
        use_aux_spatial_head=True,
    ).to(device)
    state = load_checkpoint(args.checkpoint, device)
    missing, unexpected = model.load_state_dict(state, strict=False)
    model.eval()

    x = torch.randn(1, 16, args.height, args.width, device=device)
    with torch.no_grad():
        logits = model(x)
        prob = torch.sigmoid(logits)

    expected = (1, 1, args.height, args.width)
    if tuple(prob.shape) != expected:
        raise RuntimeError(f"Expected output shape {expected}, got {tuple(prob.shape)}")
    if not torch.isfinite(prob).all():
        raise RuntimeError("Output contains non-finite values.")
    if float(prob.min()) < 0.0 or float(prob.max()) > 1.0:
        raise RuntimeError("Sigmoid output is outside [0, 1].")

    print("FireWx-FM smoke inference passed.")
    print(f"checkpoint={args.checkpoint}")
    print(f"device={device}")
    print(f"output_shape={tuple(prob.shape)}")
    print(f"prob_mean={float(prob.mean()):.6f} prob_max={float(prob.max()):.6f}")
    if missing:
        print(f"missing_keys={len(missing)}")
    if unexpected:
        print(f"unexpected_keys={len(unexpected)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())