app.py

import os
import ssl
import warnings

import spaces
import certifi

from aifs.plot import plot_field

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 tempfile

from aifs.device import  device_label

import shutil

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


PLOTTABLE_FIELDS = [
    "100u", "100v", "10u", "10v", "2t", "msl", "sp", "tcw", "swh", "mwp",
    "t_850", "t_500", "u_850", "v_850", "z_500", "q_700",
]

UNITS_MAP = {
    "2t": "K", "t_850": "K", "t_500": "K",
    "msl": "Pa", "sp": "Pa", "z_500": "m²/s²",
    "100u": "m/s", "100v": "m/s", "10u": "m/s", "10v": "m/s",
    "u_850": "m/s", "v_850": "m/s",
    "swh": "m", "mwp": "s", "tcw": "kg/m²", "q_700": "kg/kg",
}

# ── Forecast runner ───────────────────────────────────────────────────────────

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

def _phase(text: str) -> str:
    return f"**{text}**"


@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 run_forecast(lead_time: int, num_chunks: int):
    from aifs.initial_conditions import load_ics

    def emit(status, phase, dd=gr.update(choices=[]), states=[], btn=_BTN_RUNNING):
        return status, _phase(phase), dd, states, btn

    yield emit("Starting up…", "📥 Downloading initial conditions from ECMWF…")

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

        label = device_label()
        total_steps = lead_time // 6
        yield emit(
            f"Data ready — {date}",
            f"🤖 Running {lead_time}h forecast ({total_steps} steps) on {label}…",
        )

        states = []
        for kind, payload in _run_forecast_gpu(fields, date, lead_time, num_chunks):
            if kind == "log":
                yield emit(payload, f"🤖 Running inference on {label}…", states=states)
            else:
                states = payload

        timestamps = [str(s["date"]) for s in states]
        yield (
            f"✅  Forecast complete — {len(states)} steps ({lead_time}h)  ·  {label}",
            _phase(f"✅ Done — {len(states)} forecast steps ready to plot"),
            gr.update(choices=timestamps, value=timestamps[-1]),
            states,
            _BTN_READY,
        )

    except Exception as exc:
        yield (
            f"❌  Error: {exc}",
            _phase("❌ Something went wrong — see the log above"),
            gr.update(choices=[]),
            [],
            _BTN_READY,
        )

# ── Plot ──────────────────────────────────────────────────────────────────────

def plot_forecast_handler(field_name: str, timestamp: str, states: list):
    if not states:
        return None, "Run the forecast first."

    state = next(
        (s for s in states if str(s["date"]) == timestamp),
        states[-1],
    )

    if field_name not in state["fields"]:
        return None, f"Field '{field_name}' not available in this forecast state."

    values = state["fields"][field_name]
    units  = UNITS_MAP.get(field_name, "")

    fig  = plot_field(state=state, variable=field_name)
    path = os.path.join(tempfile.gettempdir(), "aifs_plot.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")

    stats = (
        f"**{field_name}** at {timestamp}\n\n"
        f"- Min: `{values.min():.4g}` {units}\n"
        f"- Max: `{values.max():.4g}` {units}\n"
        f"- Mean: `{values.mean():.4g}` {units}\n"
        f"- Grid points: `{len(values):,}`"
    )
    return path, 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; }
#phase-box { background: #161b22!important; border: 1px solid #30363d!important; border-radius: 6px; padding: 10px 14px!important; margin-bottom: 8px; }
#phase-box strong { color: #e3b341!important; font-size: 0.95em; }

/* === 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="AIFS Single v2 Forecast") as demo:
    gr.Markdown(
        """
# 🌍 AIFS Single v2 — Weather Forecast
**ECMWF's AI Integrated Forecasting System** | Runs on CUDA · Apple MPS · CPU
        """
    )

    forecast_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1, elem_classes="panel"):
            gr.Markdown("### ⚙️ Forecast Settings")
            lead_time_sl = gr.Slider(
                minimum=6, maximum=96, step=6, value=6,
                label="Lead time (hours)",
                info="On CPU/MPS keep this at 6h for testing — each step is slow without a CUDA GPU.",
            )
            num_chunks_sl = gr.Slider(
                minimum=1, maximum=32, step=1, value=16,
                label="Memory chunks",
                info="Higher = less memory, slightly slower. Ignored on CPU.",
            )
            run_btn = gr.Button("▶ Run Forecast", variant="primary", size="lg")
            phase_md = gr.Markdown(
                "Ready — click **▶ Run Forecast** to begin.",
                elem_id="phase-box",
            )
            status_box = gr.Textbox(
                label="Detailed log", lines=6, interactive=False,
                placeholder="Progress details will appear here…",
            )

        with gr.Column(scale=2, elem_classes="panel"):
            gr.Markdown("### 🗺️ Visualise Output")
            with gr.Row():
                field_dd = gr.Dropdown(
                    choices=PLOTTABLE_FIELDS, value="2t",
                    label="Field", info="Select the variable to plot.",
                )
                timestamp_dd = gr.Dropdown(
                    choices=[], label="Forecast step",
                    info="Available after running the forecast.",
                )
            plot_btn = gr.Button("🖼 Plot Field", variant="secondary")
            map_img  = gr.Image(label="Global Map", type="filepath")
            stats_md = gr.Markdown()

    run_btn.click(
        fn=run_forecast,
        inputs=[lead_time_sl, num_chunks_sl],
        outputs=[status_box, phase_md, timestamp_dd, forecast_state, run_btn],
    )
    plot_btn.click(
        fn=plot_forecast_handler,
        inputs=[field_dd, timestamp_dd, forecast_state],
        outputs=[map_img, stats_md],
    )

    gr.Markdown(
        """
---
**Notes**
- No flash-attn required — uses PyTorch SDPA (works on CPU, MPS, and CUDA).
- First run downloads the ~2 GB model checkpoint from Hugging Face and caches it.
- Data: ECMWF Open Data.
        """
    )

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