import datetime
import os
import ssl
import warnings

import spaces
import certifi

os.environ["SSL_CERT_FILE"] = certifi.where()
os.environ["REQUESTS_CA_BUNDLE"] = certifi.where()
ssl._create_default_https_context = ssl.create_default_context

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)

import gradio as gr
import numpy as np

import shutil
import tempfile

from aifs.device import device_label
from aifs.initial_conditions import EARLIEST_HISTORICAL_DATE, FULL_FIELD_RUN_HOURS
from aifs.compare import MODEL_AIFS, MODEL_WN2, MODEL_CLIMATOLOGY, MODELS, CANONICAL_FIELDS

# Clear corrupted IC cache on startup
if os.path.exists("ic_cache"):
    shutil.rmtree("ic_cache")
    os.makedirs("ic_cache", exist_ok=True)

MAX_STEPS = 8  # 48h — WeatherNext2's ~96s/step on CPU makes longer rollouts slow for a live demo
STEP_CHOICES = [str(i) for i in range(1, MAX_STEPS + 1)]

# ── AIFS GPU wrapper (unchanged from before) ───────────────────────────────────

@spaces.GPU
def _run_forecast_gpu(fields, date, lead_time, num_chunks):
    """Thin GPU-scoped wrapper — only the inference step runs inside the
    ZeroGPU allocation, so slow/retriable ECMWF downloads never eat into
    (or blow) the GPU duration budget."""
    from aifs.forecast import run_forecast as _run_forecast
    yield from _run_forecast(fields, date, lead_time=lead_time, num_chunks=num_chunks)


def _parse_historical_date(date_str: str, hour_str: str) -> datetime.datetime:
    """Parse and validate the historical-date UI inputs, or raise ValueError."""
    try:
        year, month, day = (int(p) for p in date_str.strip().split("-"))
        picked = datetime.date(year, month, day)
    except Exception:
        raise ValueError(f"could not parse date '{date_str}' — use YYYY-MM-DD.")

    if picked < EARLIEST_HISTORICAL_DATE or picked > datetime.date.today():
        raise ValueError(
            f"date must be between {EARLIEST_HISTORICAL_DATE.isoformat()} and today, got {picked}."
        )
    return datetime.datetime(picked.year, picked.month, picked.day, int(hour_str))


# ── Per-model run generators — each yields ("log", str) then ("result", states) ─

def _run_aifs(num_steps: int, historical_dt, num_chunks: int, log):
    from aifs.initial_conditions import load_ics

    yield "log", log("📥  Downloading initial conditions from ECMWF…")
    fields = date = None
    for kind, payload in load_ics(cache_dir="ic_cache", date=historical_dt):
        if kind == "log":
            yield "log", log(payload)
        else:
            fields, date = payload

    label = device_label()
    lead_time = num_steps * 6
    yield "log", log(f"🤖  Running {lead_time}h forecast ({num_steps} steps) on {label}…")

    states = []
    for kind, payload in _run_forecast_gpu(fields, date, lead_time, num_chunks):
        if kind == "log":
            yield "log", log(payload)
        else:
            states = payload

    yield "result", states


def _run_wn2(num_steps: int, log):
    from aifs import weathernext2 as wn2

    yield "log", log("📥  Fetching initial conditions from ECMWF (latest run)…")
    state = date = None
    for kind, payload in wn2.load_ics():
        if kind == "log":
            yield "log", log(payload)
        else:
            state, date = payload

    yield "log", log("🤖  Loading WeatherNext2 and running inference "
                      "(CPU — first run downloads weights + builds the mesh, then ~96s/step)…")
    states = None
    for kind, payload in wn2.run_forecast(state, date, num_steps=num_steps, device="cpu"):
        if kind == "log":
            yield "log", log(payload)
        else:
            states = payload

    yield "result", states


def _run_climatology(num_steps: int, num_years: int, log):
    from ecmwf.opendata import Client as OpendataClient

    from aifs.era5_verify import run_climatology_baseline

    # Anchor to the same "latest" reference the live models use, so steps
    # line up for comparison against an AIFS/WeatherNext2 run from "Latest".
    date = OpendataClient("ecmwf").latest()
    yield "log", log(f"📅  Anchoring climatology baseline to {date} UTC…")

    states = run_climatology_baseline(
        date, num_steps, num_years=num_years,
        log=lambda msg: log(msg),
    )
    yield "result", states


_RUN_BTN_RUNNING = gr.update(interactive=False, value="⏳ Running…")
_RUN_BTN_READY   = gr.update(interactive=True,  value="▶ Run Forecast")


def run_selected_model(
    model: str, num_steps_str: str,
    ic_mode: str, hist_date_str: str, hist_hour: str, num_chunks: int,
    clim_num_years: int,
    all_states: dict,
):
    """
    Runs whichever model is selected for `num_steps` 6h steps, storing the
    resulting states into `all_states[model]` — a dict shared across all
    three models, so a previously-run model stays available to compare
    against without re-running it.
    """
    num_steps = int(num_steps_str)
    log_lines: list[str] = []

    def log(msg: str) -> str:
        log_lines.append(msg)
        return "\n".join(log_lines)

    def emit(status, states_dict=None, btn=_RUN_BTN_RUNNING):
        return status, states_dict if states_dict is not None else all_states, btn

    try:
        if model == MODEL_AIFS:
            historical_dt = None
            if ic_mode == "Historical date":
                try:
                    historical_dt = _parse_historical_date(hist_date_str, hist_hour)
                except ValueError as exc:
                    yield emit(f"❌  {exc}", btn=_RUN_BTN_READY)
                    return
            runner = _run_aifs(num_steps, historical_dt, int(num_chunks), log)
        elif model == MODEL_WN2:
            runner = _run_wn2(num_steps, log)
        else:
            runner = _run_climatology(num_steps, int(clim_num_years), log)

        states = None
        for kind, payload in runner:
            if kind == "log":
                yield emit(payload)
            else:
                states = payload

        new_all_states = dict(all_states)
        new_all_states[model] = states
        yield emit(log(f"✅  {model} — {len(states)} step(s) ready to plot/compare."), new_all_states, _RUN_BTN_READY)

    except Exception as exc:
        yield emit(log(f"❌  Error: {exc}"), btn=_RUN_BTN_READY)


def toggle_model_controls(model: str):
    return (
        gr.update(visible=(model == MODEL_AIFS)),
        gr.update(visible=(model == MODEL_CLIMATOLOGY)),
    )


def format_status(all_states: dict) -> str:
    """One line per model: whether it's been run this session and, if so, its step range."""
    lines = []
    for model in MODELS:
        states = all_states.get(model)
        if states:
            lines.append(f"- **{model}**: ✅ {len(states)} step(s) ready — {states[0]['date']} → {states[-1]['date']}")
        else:
            lines.append(f"- **{model}**: ⬜ not run yet")
    return "\n".join(lines)


# ── Forecast archive (save/load runs to the HF dataset in aifs.archive) ───────

def save_current_run(model: str, all_states: dict) -> str:
    from aifs.archive import save_run

    states = all_states.get(model)
    if not states:
        return f"⚠️  Run {model} first, then save it."

    log_lines: list[str] = []
    try:
        save_run(model, states, log=log_lines.append)
        return "\n".join(log_lines)
    except Exception as exc:
        return "\n".join(log_lines) + f"\n❌  Error: {exc}"


def refresh_saved_runs():
    from aifs.archive import list_saved_runs

    log_lines: list[str] = []
    try:
        runs = list_saved_runs(log=log_lines.append)
    except Exception as exc:
        return gr.update(choices=[], value=None), "\n".join(log_lines) + f"\n❌  Error: {exc}"

    choices = [
        (
            f"{r['model']} — init {r['init_date']:%Y-%m-%d %H:%M} — {r['num_steps']} step(s) "
            f"— saved {r['saved_at']:%Y-%m-%d %H:%M}",
            r["filename"],
        )
        for r in runs
    ]
    return gr.update(choices=choices, value=(choices[0][1] if choices else None)), "\n".join(log_lines)


def load_saved_run(filename: str, all_states: dict):
    from aifs.archive import load_run

    if not filename:
        return all_states, "⚠️  Pick a saved run to load first."

    log_lines: list[str] = []
    try:
        model, states = load_run(filename, log=log_lines.append)
    except Exception as exc:
        return all_states, "\n".join(log_lines) + f"\n❌  Error: {exc}"

    new_all_states = dict(all_states)
    new_all_states[model] = states
    return new_all_states, "\n".join(log_lines)


# ── Visualize (any one model/step/field) ───────────────────────────────────────

def plot_selected(model: str, step_str: str, field: str, all_states: dict):
    from aifs.compare import extract, plot_model_field

    states = all_states.get(model) or []
    if not states:
        return None, f"Run **{model}** first."

    step_idx = min(int(step_str) - 1, len(states) - 1)
    try:
        fig = plot_model_field(model, states, step_idx, field)
        _, _, values, _ = extract(model, states, step_idx, field)
    except Exception as exc:
        return None, f"❌  {exc}"

    path = os.path.join(tempfile.gettempdir(), "viz_plot.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")

    spec = CANONICAL_FIELDS[field]
    date = states[step_idx]["date"]
    stats = (
        f"**{model}** — {spec['long_name']} @ {date}  (step {step_idx + 1})\n\n"
        f"- Min: `{np.nanmin(values):.4g}` {spec['units']}\n"
        f"- Max: `{np.nanmax(values):.4g}` {spec['units']}\n"
        f"- Mean: `{np.nanmean(values):.4g}` {spec['units']}\n"
    )
    return path, stats


# ── Compare (any two models/steps, one field) ──────────────────────────────────

def compare_selected(model_a: str, step_a_str: str, model_b: str, step_b_str: str, field: str, all_states: dict):
    from aifs.compare import plot_compare_maps

    states_a = all_states.get(model_a) or []
    states_b = all_states.get(model_b) or []
    if not states_a:
        return None, None, None, f"Run **{model_a}** first."
    if not states_b:
        return None, None, None, f"Run **{model_b}** first."

    idx_a = min(int(step_a_str) - 1, len(states_a) - 1)
    idx_b = min(int(step_b_str) - 1, len(states_b) - 1)

    try:
        fig_a, fig_b, fig_diff, result = plot_compare_maps(model_a, states_a, idx_a, model_b, states_b, idx_b, field)
    except Exception as exc:
        return None, None, None, f"❌  {exc}"

    paths = []
    for name, fig in zip(("a", "b", "diff"), (fig_a, fig_b, fig_diff)):
        path = os.path.join(tempfile.gettempdir(), f"cmp_{name}.png")
        fig.savefig(path, dpi=150, bbox_inches="tight")
        paths.append(path)

    spec = CANONICAL_FIELDS[field]
    stats = (
        f"**{spec['long_name']}** — {model_a} (step {idx_a + 1}) vs {model_b} (step {idx_b + 1})\n\n"
        f"| Metric | Value |\n|---|---|\n"
        f"| RMSE | `{result['rmse']:.4g}` {spec['units']} |\n"
        f"| MAE | `{result['mae']:.4g}` {spec['units']} |\n"
        f"| Bias (A − B) | `{result['bias']:.4g}` {spec['units']} |\n"
        f"| Correlation | `{result['corr']:.3f}` |\n"
        f"| Points compared | `{result['n']:,}` |\n"
    )
    return paths[0], paths[1], paths[2], stats


# ── UI ────────────────────────────────────────────────────────────────────────
DARK_CSS = """
body, .gradio-container {
    background: #0d1117!important; color: #cdd9e5!important;
    font-family: 'Inter','Segoe UI',sans-serif;
}
h1 { color: #58a6ff!important; letter-spacing: -0.5px; }
h3 { color: #79c0ff!important; }
.panel { background: #161b22!important; border: 1px solid #30363d!important; border-radius: 8px; }
button.primary { background: #1f6feb!important; border: none!important; color: white!important; }
button.primary:hover { background: #388bfd!important; }
button.primary:disabled { background: #30363d!important; color: #8b949e!important; cursor: not-allowed!important; }
.label-wrap { color: #8b949e!important; }
textarea, input, select { background: #1c2128!important; color: #cdd9e5!important; border-color: #30363d!important; }
textarea::placeholder, input::placeholder { color: #a8b3c0!important; }
.output-markdown { color: #cdd9e5!important; }

/* === Dropdown / radio internals ===
   Gradio's Dropdown/Radio are custom components, not a plain <select>/
   <input> — the rule above doesn't reach their popup list or selected-pill
   styling, which otherwise falls back to Gradio's default (light) theme. */
ul[class*="options"], li[class*="item"] {
    background: #1c2128!important;
    color: #cdd9e5!important;
}
li[class*="item"]:hover, li[class*="item"][aria-selected="true"] {
    background: #30363d!important;
    color: #cdd9e5!important;
}
.wrap label {
    background: #1c2128!important;
    color: #cdd9e5!important;
    border-color: #30363d!important;
}
label.selected, label[class*="selected"] {
    background: #1f6feb!important;
    color: #ffffff!important;
    border-color: #1f6feb!important;
}

/* === Notes / Markdown prose === */
.prose,
.prose p,
.prose li,
.prose ul,
.prose ol,
.prose strong {
    color: #a8b3c0!important;
}
.prose h1,
.prose h2,
.prose h3,
.prose h4 {
    color: #a8b3c0!important;
}

footer { display: none!important; }
"""

with gr.Blocks(css=DARK_CSS, title="Weather Model Comparison") as demo:
    gr.Markdown(
        """
# 🌍 Weather Model Comparison
**AIFS Single v2 (ECMWF)** · **WeatherNext 2 (Google DeepMind)** · **ERA5 climatology baseline**
        """
    )

    all_states = gr.State({})

    with gr.Group(elem_classes="panel"):
        gr.Markdown("### 📋 Session Status")
        status_md = gr.Markdown(format_status({}))

    with gr.Row():
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### ⚙️ Run a Forecast")
            model_dd = gr.Radio(
                MODELS, value=MODEL_AIFS, label="Model",
                info="AIFS runs on GPU (ZeroGPU); WeatherNext2 runs on CPU (~96s/step — "
                     "needs ~50GB RAM/first-load, more than this Space's default GPU "
                     "allocation); Climatology is a zero-skill baseline for comparison, not a real model.",
            )
            num_steps_dd = gr.Dropdown(
                STEP_CHOICES, value="2", label="Number of steps (6h each)",
            )

            with gr.Group(visible=True) as aifs_controls:
                ic_mode_radio = gr.Radio(
                    ["Latest", "Historical date"], value="Latest",
                    label="Initial conditions (AIFS only)",
                    info="Historical dates pull from ECMWF's deeper S3 archive (from "
                         f"{EARLIEST_HISTORICAL_DATE.isoformat()}) instead of the live feed.",
                )
                with gr.Row(visible=False) as historical_row:
                    hist_date_tb = gr.Textbox(
                        label="Date (UTC)", placeholder="YYYY-MM-DD",
                        value=(datetime.date.today() - datetime.timedelta(days=7)).isoformat(),
                    )
                    hist_hour_dd = gr.Dropdown(
                        [f"{h:02d}" for h in FULL_FIELD_RUN_HOURS], value="00",
                        label="Run hour (UTC)",
                        info="Limited to 00/12 UTC — 06/18 UTC used a reduced ECMWF product "
                             "before 2026-05-12 that's missing fields AIFS needs.",
                    )
                ic_mode_radio.change(
                    fn=lambda mode: gr.update(visible=(mode == "Historical date")),
                    inputs=ic_mode_radio, outputs=historical_row,
                )
                num_chunks_sl = gr.Slider(
                    minimum=1, maximum=32, step=1, value=16,
                    label="Memory chunks (AIFS only)",
                    info="Higher = less memory, slightly slower. Ignored on CPU.",
                )

            with gr.Group(visible=False) as clim_controls:
                clim_years_sl = gr.Slider(
                    minimum=3, maximum=30, step=1, value=10,
                    label="Climatology years (baseline only)",
                    info="How many past years of ERA5 to average per step.",
                )

            model_dd.change(
                fn=toggle_model_controls, inputs=model_dd, outputs=[aifs_controls, clim_controls],
            )

            run_btn = gr.Button("▶ Run Forecast", variant="primary", size="lg")
            run_status = gr.Textbox(
                label="Detailed log", lines=8, interactive=False,
                placeholder="Progress details will appear here…",
            )
            save_btn = gr.Button("💾 Save Current Run to Archive", variant="secondary")
            save_status = gr.Textbox(label="Save log", lines=2, interactive=False)

        with gr.Column(scale=2, elem_classes="panel"):
            gr.Markdown("### 🗺️ Visualize")
            with gr.Row():
                viz_model_dd = gr.Dropdown(MODELS, value=MODEL_AIFS, label="Model")
                viz_step_dd = gr.Dropdown(STEP_CHOICES, value="1", label="Step")
                viz_field_dd = gr.Dropdown(
                    sorted(CANONICAL_FIELDS), value="2m_temperature", label="Field",
                )
            viz_btn = gr.Button("🖼 Plot Field", variant="secondary")
            viz_img = gr.Image(label="Map", type="filepath")
            viz_stats_md = gr.Markdown()

    gr.Markdown("---")
    with gr.Row():
        with gr.Column(elem_classes="panel"):
            gr.Markdown(
                "### 📂 Saved Runs\n"
                "Forecasts saved above persist in the "
                "[weather-forecast-archive](https://huggingface.co/datasets/EmmaScharfmann/weather-forecast-archive) "
                "dataset — load one back here to plot/compare it without re-running the model."
            )
            with gr.Row():
                saved_runs_dd = gr.Dropdown(choices=[], label="Saved run", scale=3)
                refresh_saved_btn = gr.Button("🔄 Refresh", scale=1)
                load_saved_btn = gr.Button("📥 Load", variant="secondary", scale=1)
            saved_runs_status = gr.Textbox(label="Archive log", lines=2, interactive=False)

    gr.Markdown("---")
    with gr.Row():
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown(
                "### 📊 Compare Two Models\n"
                "Pick any two model runs (including two steps of the same model) and a "
                "field — shows both maps, their difference, and a skill metric (RMSE, "
                "MAE, bias, correlation). AIFS's irregular grid is compared by sampling "
                "the other model onto AIFS's own points; WeatherNext2 and the climatology "
                "baseline share an identical grid, so no resampling is needed between them."
            )
            with gr.Row():
                cmp_model_a_dd = gr.Dropdown(MODELS, value=MODEL_AIFS, label="Model A")
                cmp_step_a_dd = gr.Dropdown(STEP_CHOICES, value="1", label="Step A")
            with gr.Row():
                cmp_model_b_dd = gr.Dropdown(MODELS, value=MODEL_CLIMATOLOGY, label="Model B")
                cmp_step_b_dd = gr.Dropdown(STEP_CHOICES, value="1", label="Step B")
            cmp_field_dd = gr.Dropdown(
                sorted(CANONICAL_FIELDS), value="2m_temperature", label="Field",
            )
            cmp_btn = gr.Button("📊 Compare", variant="primary")

        with gr.Column(scale=2, elem_classes="panel"):
            with gr.Row():
                cmp_img_a = gr.Image(label="Model A", type="filepath")
                cmp_img_b = gr.Image(label="Model B", type="filepath")
                cmp_img_diff = gr.Image(label="A − B", type="filepath")
            cmp_stats_md = gr.Markdown()

    run_btn.click(
        fn=run_selected_model,
        inputs=[model_dd, num_steps_dd, ic_mode_radio, hist_date_tb, hist_hour_dd, num_chunks_sl,
                clim_years_sl, all_states],
        outputs=[run_status, all_states, run_btn],
    )
    viz_btn.click(
        fn=plot_selected,
        inputs=[viz_model_dd, viz_step_dd, viz_field_dd, all_states],
        outputs=[viz_img, viz_stats_md],
    )
    cmp_btn.click(
        fn=compare_selected,
        inputs=[cmp_model_a_dd, cmp_step_a_dd, cmp_model_b_dd, cmp_step_b_dd, cmp_field_dd, all_states],
        outputs=[cmp_img_a, cmp_img_b, cmp_img_diff, cmp_stats_md],
    )
    save_btn.click(
        fn=save_current_run,
        inputs=[model_dd, all_states],
        outputs=[save_status],
    )
    refresh_saved_btn.click(
        fn=refresh_saved_runs,
        inputs=[],
        outputs=[saved_runs_dd, saved_runs_status],
    )
    load_saved_btn.click(
        fn=load_saved_run,
        inputs=[saved_runs_dd, all_states],
        outputs=[all_states, saved_runs_status],
    )
    all_states.change(
        fn=format_status,
        inputs=all_states,
        outputs=status_md,
    )

    gr.Markdown(
        """
---
**Notes**
- AIFS: no flash-attn required (PyTorch SDPA — works on CPU, MPS, CUDA). First run downloads its ~2GB checkpoint.
- WeatherNext2: runs via an unmerged `transformers` fork; CPU-only here, ~96s/step once loaded.
- Climatology baseline: pure ERA5 climatological mean per step — zero model skill by construction, a reference point for judging whether AIFS/WeatherNext2 add value.
- Data: ECMWF Open Data (forecasts) and EarthMover's public ERA5 archive (climatology).
        """
    )

if __name__ == "__main__":
    demo.launch()