"""Extract a single (lat, lon) point's time series from a saved global AIFS
forecast (produced by aifs_run_global.py) and write a CSV. No GPU needed -
just an array lookup, so this runs in well under a second.

Pipeline step 2/2: 都市名 -> 緯度経度 -> Grid Point検索 -> Pickup(この保存済み全球データから) -> CSV作成

Usage:
    python extract_city_from_global.py --npz global_forecast.npz --lat 34.3853 --lon 132.4553 --out city.csv
"""

import argparse
import csv
import datetime
import math

import numpy as np


def rh_percent(t_c, td_c):
    a, b = 17.625, 243.04
    es = math.exp((a * t_c) / (b + t_c))
    e = math.exp((a * td_c) / (b + td_c))
    return round(100.0 * e / es, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--npz", type=str, required=True)
    ap.add_argument("--lat", type=float, required=True)
    ap.add_argument("--lon", type=float, required=True)
    ap.add_argument("--out", type=str, required=True)
    args = ap.parse_args()

    data = np.load(args.npz)
    lat_arr = data["latitudes"]
    lon_arr = data["longitudes"]
    dates = data["dates"]

    lon_signed = np.where(lon_arr > 180, lon_arr - 360, lon_arr)
    dist2 = (lat_arr - args.lat) ** 2 + (lon_signed - args.lon) ** 2
    idx = int(np.argmin(dist2))
    print(f"Nearest grid point: lat={lat_arr[idx]:.3f} lon={lon_arr[idx]:.3f} (idx={idx})")

    n_steps = len(dates)
    rows = []
    for i in range(n_steps):
        t2m_c = float(data["2t"][i, idx]) - 273.15
        d2m_c = float(data["2d"][i, idx]) - 273.15
        u10 = float(data["10u"][i, idx])
        v10 = float(data["10v"][i, idx])
        wind_speed = float(np.hypot(u10, v10))
        msl_hpa = float(data["msl"][i, idx]) / 100.0
        tcc = float(data["tcc"][i, idx])
        tp_m = float(data["tp"][i, idx])
        precip_6h_mm = tp_m * 1000.0

        dt = datetime.datetime.fromisoformat(str(dates[i]))
        rows.append(
            dict(
                date=dt.isoformat(),
                date_jst=(dt + datetime.timedelta(hours=9)).strftime("%Y-%m-%d %H:%M") + " JST",
                t2m_c=round(t2m_c, 1),
                humidity_pct=rh_percent(t2m_c, d2m_c),
                wind_speed_ms=round(wind_speed, 1),
                msl_hpa=round(msl_hpa, 1),
                cloud_cover=round(tcc, 2),
                precip_6h_mm=round(precip_6h_mm, 2),
            )
        )

    with open(args.out, "w", newline="", encoding="utf-8") as fp:
        writer = csv.DictWriter(fp, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)
    print(f"DONE: wrote {len(rows)} rows to {args.out}")


if __name__ == "__main__":
    main()