"""Streamlit GUI for the AIFS-Single 15-day weather trial.

Reads a single saved global forecast (aifs_run_global.py output) and lets
you pick ANY location - extraction is a plain array lookup (no GPU needed).
The "全球データを計算する" button is the only thing that touches the GPU.

Run with:
    <aifs conda env python> -m streamlit run aifs_weather_app.py
"""

import datetime
import math
import os
import pathlib
import subprocess
import sys

import numpy as np
import plotly.graph_objects as go
import requests
import streamlit as st

BASE = pathlib.Path(r"<YOUR_PROJECT_ROOT>")
GLOBAL_SCRIPT = BASE / "tools" / "aifs_run_global.py"
MAP_SCRIPT = BASE / "tools" / "draw_weather_map.py"
FORECAST_DIR = BASE / "experiments" / "forecasts"
NPZ_PATH = FORECAST_DIR / "global_forecast.npz"
MAP_CACHE_DIR = FORECAST_DIR / "maps"
PYTHON_EXE = pathlib.Path(sys.executable)

# conda env の python.exe を直接叩くとPATHにLibrary/binが入らず、
# matplotlib/cartopyの描画がネイティブクラッシュすることがある対策
CONDA_ENV_DIR = PYTHON_EXE.parent
_EXTRA_DLL_PATHS = [
    str(CONDA_ENV_DIR),
    str(CONDA_ENV_DIR / "Library" / "mingw-w64" / "bin"),
    str(CONDA_ENV_DIR / "Library" / "usr" / "bin"),
    str(CONDA_ENV_DIR / "Library" / "bin"),
    str(CONDA_ENV_DIR / "Scripts"),
]


def subprocess_env():
    env = os.environ.copy()
    env["PYTHONIOENCODING"] = "utf-8"
    env["KMP_DUPLICATE_LIB_OK"] = "TRUE"
    env["PATH"] = os.pathsep.join(_EXTRA_DLL_PATHS) + os.pathsep + env.get("PATH", "")
    return env


def generate_map(hours):
    MAP_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    out_path = MAP_CACHE_DIR / f"map_{hours}h.png"
    cmd = [str(PYTHON_EXE), str(MAP_SCRIPT), "--hours", str(hours), "--out", str(out_path)]
    proc = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", env=subprocess_env())
    return out_path, proc.returncode, proc.stdout + proc.stderr

CITIES = {
    "広島市": (34.3853, 132.4553),
    "東京都": (35.6812, 139.7671),
    "大阪市": (34.6937, 135.5023),
    "札幌市": (43.0621, 141.3544),
    "那覇市": (26.2124, 127.6809),
    "カスタム": None,
}


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)


@st.cache_data
def geocode_city(name):
    """City name -> (lat, lon, display_name), or None if not found.
    Uses OpenStreetMap's free Nominatim API (no key required)."""
    try:
        resp = requests.get(
            "https://nominatim.openstreetmap.org/search",
            params={"q": name, "format": "json", "limit": 1, "accept-language": "ja"},
            headers={"User-Agent": "aifs-weather-app/1.0 (personal local trial)"},
            timeout=10,
        )
        resp.raise_for_status()
        results = resp.json()
        if not results:
            return None
        r = results[0]
        return float(r["lat"]), float(r["lon"]), r["display_name"]
    except Exception as e:
        return "error", str(e), None


@st.cache_data
def load_global(npz_path, mtime):
    # mtime is passed just to bust the cache when the file is refreshed
    return dict(np.load(npz_path))


def extract_point(data, lat, lon):
    lat_arr = data["latitudes"]
    lon_arr = data["longitudes"]
    lon_signed = np.where(lon_arr > 180, lon_arr - 360, lon_arr)
    dist2 = (lat_arr - lat) ** 2 + (lon_signed - lon) ** 2
    idx = int(np.argmin(dist2))

    dates = data["dates"]
    rows = []
    for i in range(len(dates)):
        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])
        precip_6h_mm = float(data["tp"][i, idx]) * 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),
            )
        )
    return rows, lat_arr[idx], lon_arr[idx]


st.set_page_config(page_title="AIFS 15日天気トライアル", page_icon="\U0001F324", layout="wide")
st.title("AIFS-Single 15日天気予報トライアル")

st.warning(
    "**この予報の精度には既知の限界があります。** 検証の結果、猛暑のような極端な気象では"
    "2〜3日目以降に実際の予報と大きく乖離することを確認しています"
    "（MSE損失で学習された軽量モデルの既知の弱点の可能性）。実用の天気予報としては使わないでください。"
)

col1, col2 = st.columns([1, 2])

with col1:
    st.subheader("① 全球データ（1回だけGPU計算）")
    if NPZ_PATH.exists():
        mtime = datetime.datetime.fromtimestamp(NPZ_PATH.stat().st_mtime)
        size_mb = NPZ_PATH.stat().st_size / 1e6
        st.caption(f"計算済み: {mtime.strftime('%Y-%m-%d %H:%M')}（{size_mb:.0f} MB）")
    else:
        st.caption("まだ計算されていません")

    run_clicked = st.button("全球データを計算する（GPU計算、約15〜20分）", type="primary")

    if run_clicked:
        status = st.status("全球データを計算中... (実データ取得 → GPU推論 → 保存)", expanded=True)
        cmd = [str(PYTHON_EXE), str(GLOBAL_SCRIPT), "--lead-hours", "360", "--out", str(NPZ_PATH)]
        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", env=subprocess_env()
        )
        log_lines = []
        for line in proc.stdout:
            log_lines.append(line.rstrip())
            status.write(line.rstrip())
        proc.wait()
        if proc.returncode == 0:
            status.update(label="完了しました", state="complete")
            st.cache_data.clear()
        else:
            status.update(label=f"エラーで終了しました（コード {proc.returncode}）", state="error")
            st.code("\n".join(log_lines[-30:]))

    st.divider()
    st.subheader("② 地点を選ぶ（GPU不要、瞬時に抽出）")
    city = st.selectbox("都市を選択", list(CITIES.keys()))
    if city == "カスタム":
        place_name = st.text_input("都市名・地名を入力", value="京都市")
        lat = lon = None
        if place_name:
            result = geocode_city(place_name)
            if result is None:
                st.error("場所が見つかりませんでした。別の書き方（例: 「Kyoto, Japan」）を試してください。")
            elif result[0] == "error":
                st.error(f"検索エラー: {result[1]}")
            else:
                lat, lon, display_name = result
                st.success(f"推定位置: {display_name}\n\n緯度={lat:.4f}, 経度={lon:.4f}")
        with st.expander("緯度経度を直接指定する"):
            manual = st.checkbox("手入力を使う")
            if manual:
                lat = st.number_input("緯度", value=lat or 35.0, format="%.4f")
                lon = st.number_input("経度", value=lon or 135.0, format="%.4f")
    else:
        lat, lon = CITIES[city]
        st.write(f"緯度: {lat}, 経度: {lon}")

if NPZ_PATH.exists() and lat is not None and lon is not None:
    global_data = load_global(str(NPZ_PATH), NPZ_PATH.stat().st_mtime)
    rows, grid_lat, grid_lon = extract_point(global_data, lat, lon)

    with col2:
        st.subheader(f"{city} の予報")
        st.caption(f"最寄り格子点: lat={grid_lat:.3f}, lon={grid_lon:.3f}")

        x = [r["date_jst"] for r in rows]

        fig_temp = go.Figure()
        fig_temp.add_trace(go.Scatter(x=x, y=[r["t2m_c"] for r in rows], mode="lines+markers", name="気温(℃)"))
        fig_temp.update_layout(title="気温", xaxis_title="日時(JST)", yaxis_title="℃", height=300)
        st.plotly_chart(fig_temp, use_container_width=True)

        fig_hum = go.Figure()
        fig_hum.add_trace(
            go.Scatter(x=x, y=[r["humidity_pct"] for r in rows], mode="lines+markers", name="湿度(%)", line=dict(color="green"))
        )
        fig_hum.update_layout(title="相対湿度", xaxis_title="日時(JST)", yaxis_title="%", height=300)
        st.plotly_chart(fig_hum, use_container_width=True)

        fig_precip = go.Figure()
        fig_precip.add_trace(go.Bar(x=x, y=[r["precip_6h_mm"] for r in rows], name="6時間降水量(mm)"))
        fig_precip.update_layout(title="6時間降水量", xaxis_title="日時(JST)", yaxis_title="mm", height=300)
        st.plotly_chart(fig_precip, use_container_width=True)

        fig_wind = go.Figure()
        fig_wind.add_trace(
            go.Scatter(x=x, y=[r["wind_speed_ms"] for r in rows], mode="lines+markers", name="風速(m/s)", line=dict(color="orange"))
        )
        fig_wind.update_layout(title="風速", xaxis_title="日時(JST)", yaxis_title="m/s", height=300)
        st.plotly_chart(fig_wind, use_container_width=True)

        st.dataframe(rows, use_container_width=True)
else:
    with col2:
        st.info("まず左のパネルで「全球データを計算する」を実行してください（初回のみGPU計算が必要です）。")

if NPZ_PATH.exists():
    st.divider()
    st.subheader("③ 天気図（日本周辺、GPU不要）")

    map_col1, map_col2 = st.columns(2)

    with map_col1:
        st.markdown("**今（0時間後）**")
        out_path, code, log = generate_map(0)
        if code == 0 and out_path.exists():
            st.image(str(out_path), use_container_width=True)
        else:
            st.error("天気図の生成に失敗しました")
            st.code(log[-2000:])

    with map_col2:
        lead_hours_map = st.slider("何時間後の天気図を見るか", min_value=0, max_value=360, value=24, step=6, key="map_hours")
        st.markdown(f"**{lead_hours_map}時間後**")
        out_path2, code2, log2 = generate_map(lead_hours_map)
        if code2 == 0 and out_path2.exists():
            st.image(str(out_path2), use_container_width=True)
        else:
            st.error("天気図の生成に失敗しました")
            st.code(log2[-2000:])