#!/usr/bin/env python3
"""Extract official NeuralGCM static fields from released checkpoints."""
from __future__ import annotations

import argparse
import pickle

try:
    from common import PROJECT_ROOT, load_config, resolve_path
except ModuleNotFoundError:  # supports ``python -m scripts.prepare_static_data``
    from scripts.common import PROJECT_ROOT, load_config, resolve_path


STATIC_VARIABLES = ("geopotential_at_surface", "land_sea_mask")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
    parser.add_argument(
        "--mode",
        action="append",
        dest="modes",
        help="profile to extract; repeat for multiple profiles (default: all)",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="replace an existing profile static NetCDF",
    )
    args = parser.parse_args()
    config = load_config(args.config)
    profiles = config["model"]["profiles"]
    static_files = config["data"].get("static_files", {})
    modes = args.modes or list(profiles)

    import xarray as xr

    for mode in modes:
        if mode not in profiles:
            raise ValueError(f"Unknown model profile {mode!r}")
        if mode not in static_files:
            raise ValueError(f"data.static_files has no path for {mode!r}")
        checkpoint_path = resolve_path(
            profiles[mode]["official_reference"], args.config
        )
        output_path = resolve_path(static_files[mode], args.config)
        if output_path.exists() and not args.overwrite:
            print(f"Static fields already exist: {output_path}")
            continue
        with checkpoint_path.open("rb") as handle:
            payload = pickle.load(handle)
        if not isinstance(payload, dict) or "aux_ds_dict" not in payload:
            raise ValueError(
                f"{checkpoint_path} is missing the official aux_ds_dict"
            )
        source = xr.Dataset.from_dict(payload["aux_ds_dict"])
        missing = [name for name in STATIC_VARIABLES if name not in source]
        if missing:
            raise ValueError(
                f"{checkpoint_path} is missing static variables {missing}"
            )
        static = source[list(STATIC_VARIABLES)]
        static.attrs.update(
            {
                "source": str(checkpoint_path),
                "neuralgcm_profile": mode,
                "synthetic": "false",
            }
        )
        output_path.parent.mkdir(parents=True, exist_ok=True)
        static.to_netcdf(output_path)
        print(
            f"Extracted {mode} static fields to {output_path}: "
            f"longitude={static.sizes['longitude']}, "
            f"latitude={static.sizes['latitude']}"
        )


if __name__ == "__main__":
    main()