#!/usr/bin/env python3
"""Visualize official NeuralGCM pressure-level predictions."""
from __future__ import annotations
import argparse
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
try:
    from common import load_config, resolve_path
except ModuleNotFoundError:  # supports ``python -m scripts.result``
    from scripts.common import load_config, resolve_path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--input")
    parser.add_argument("--variable", default="temperature_500")
    parser.add_argument("--level", type=int, default=500)
    parser.add_argument("--lead", type=int, default=0)
    parser.add_argument("--output")
    args = parser.parse_args()
    config = load_config(args.config)
    ds = xr.open_dataset(resolve_path(args.input or config["inference"].get("output", "results/predictions.nc"), args.config))
    # Official PressureLevelModel output stores named variables with an
    # explicit pressure ``level`` dimension; retain compatibility with the
    # earlier channel-packed smoke output when present.
    variable = args.variable
    if variable not in ds.data_vars and variable.rsplit("_", 1)[-1].isdigit():
        base, suffix = variable.rsplit("_", 1)
        if base in ds.data_vars:
            variable = base
            args.level = int(suffix)
    if variable in ds.data_vars:
        field = ds[variable]
        if "level" in field.dims:
            level = int(args.level)
            if level not in ds.level.values:
                raise ValueError(f"Unknown pressure level {level}; available levels={ds.level.values.tolist()}")
            field = field.sel(level=level)
    elif "channel" in ds.coords and args.variable in [str(x) for x in ds.channel.values]:
        field = ds.prediction.sel(channel=args.variable)
    else:
        raise ValueError(f"Unknown channel/variable {args.variable!r}")
    if not bool(np.isfinite(field.values).all()):
        raise ValueError(
            f"Variable {variable!r} contains NaN/Inf; refusing to create "
            "a misleading visualization"
        )
    field.isel(time=args.lead).plot(figsize=(12, 4), cmap="coolwarm")
    output = resolve_path(args.output or "results/forecast.png", args.config)
    output.parent.mkdir(parents=True, exist_ok=True)
    plt.tight_layout(); plt.savefig(output, dpi=150)
    print(f"Saved visualization to {output}")


if __name__ == "__main__":
    main()