from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import xarray as xr

from common import load_config, resolve_path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--input", default=None)
    parser.add_argument("--channel", default=None)
    args = parser.parse_args()
    cfg = load_config(args.config)
    viz = cfg["visualization"]
    source = Path(args.input) if args.input else resolve_path(viz["input_file"], cfg)
    channel = args.channel or viz["channel"]
    data = xr.open_dataarray(source).sel(channel=channel)
    if "time" in data.dims:
        data = data.isel(time=-1)
    output = resolve_path(viz["output_file"], cfg)
    output.parent.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(12, 5), constrained_layout=True)
    image = ax.pcolormesh(data.lon, data.lat, data, shading="auto", cmap=viz["cmap"])
    valid_time = str(data.time.values) if "time" in data.coords else ""
    ax.set(title=f"FuXi 2.1 {channel} | {valid_time}", xlabel="Longitude", ylabel="Latitude")
    fig.colorbar(image, ax=ax, label=channel)
    fig.savefig(output, dpi=160)
    plt.close(fig)
    print(f"Saved visualization to {output}")


if __name__ == "__main__":
    main()