"""Draw a Japan-region weather map (temperature fill + MSL isobars + 10m
wind vectors) from a saved global AIFS forecast (aifs_run_global.py output).
No GPU needed - just reads one timestep out of the saved array.

Usage:
    python draw_weather_map.py --hours 24 --out map_24h.png
"""

import argparse
import datetime
import pathlib

import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import matplotlib.tri as tri
import numpy as np

plt.rcParams["font.family"] = "Yu Gothic"

NPZ_PATH = r"<YOUR_PROJECT_ROOT>\forecasts\global_forecast.npz"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--hours", type=int, default=0, help="Lead time in hours (multiple of 6)")
    ap.add_argument("--out", type=str, required=True)
    args = ap.parse_args()

    data = np.load(NPZ_PATH)
    lat = data["latitudes"]
    lon = data["longitudes"]
    lon_signed = np.where(lon > 180, lon - 360, lon)
    dates = data["dates"]

    step = args.hours // 6
    if step >= len(dates):
        raise ValueError(f"--hours too large: only {len(dates)} steps (0-{(len(dates)-1)*6}h) available")

    dt = datetime.datetime.fromisoformat(str(dates[step]))

    region = (lon_signed > 115) & (lon_signed < 155) & (lat > 20) & (lat < 50)

    msl_hpa = data["msl"][step] / 100.0
    t2m_c = data["2t"][step] - 273.15
    u10 = data["10u"][step]
    v10 = data["10v"][step]

    fig = plt.figure(figsize=(10, 9))
    ax = plt.axes(projection=ccrs.PlateCarree())
    ax.set_extent([115, 155, 20, 50], crs=ccrs.PlateCarree())
    ax.add_feature(cfeature.LAND, facecolor="#f0f0e8")
    ax.add_feature(cfeature.OCEAN, facecolor="#dbe9f4")
    ax.add_feature(cfeature.COASTLINE, linewidth=0.6)
    ax.add_feature(cfeature.BORDERS, linewidth=0.4, linestyle=":")

    triang = tri.Triangulation(lon_signed[region], lat[region])
    temp_fill = ax.tricontourf(triang, t2m_c[region], levels=20, cmap="RdYlBu_r", alpha=0.75, transform=ccrs.PlateCarree())
    cbar = plt.colorbar(temp_fill, ax=ax, orientation="vertical", shrink=0.7, pad=0.02)
    cbar.set_label("気温 (℃)")

    pressure_levels = np.arange(np.floor(msl_hpa[region].min() / 4) * 4, np.ceil(msl_hpa[region].max() / 4) * 4 + 4, 4)
    cs = ax.tricontour(
        triang, msl_hpa[region], levels=pressure_levels, colors="black", linewidths=0.8, transform=ccrs.PlateCarree()
    )
    ax.clabel(cs, inline=True, fontsize=7, fmt="%d")

    step_thin = 8
    ax.quiver(
        lon_signed[region][::step_thin],
        lat[region][::step_thin],
        u10[region][::step_thin],
        v10[region][::step_thin],
        transform=ccrs.PlateCarree(),
        scale=300,
        width=0.002,
        color="dimgray",
    )

    gl = ax.gridlines(draw_labels=True, linewidth=0.3, linestyle="--", alpha=0.5)
    gl.top_labels = False
    gl.right_labels = False

    jst = dt + datetime.timedelta(hours=9)
    fig.text(
        0.5,
        0.93,
        f"AIFS-Single 予報天気図（{args.hours}時間後）\n{jst.strftime('%Y-%m-%d %H:%M')} JST（等圧線=海面気圧, 色=気温, 矢印=地上10m風）",
        ha="center",
        fontsize=13,
    )

    out_path = pathlib.Path(args.out)
    plt.savefig(out_path, dpi=150)
    print("Saved:", out_path)


if __name__ == "__main__":
    main()