"""Gradio app — primary deployment target for HF Spaces.

Generates 24-hour scenarios with PC-DDPM and overlays them on a chosen
reference day from the demo bank. Inference runs live (T=1000 reverse
diffusion on CPU); a progress bar manages the wait per the
2026-05-02 inference-is-live decision. The GNN surrogate then scores
voltage feasibility under operational and ANSI bounds — the result
shows up as a traffic-light badge next to the plot.
"""

from __future__ import annotations

import gradio as gr
import matplotlib.figure

from pc_ddpm_alberta.feasibility import (
    V_BOUNDS_ANSI,
    V_BOUNDS_OPERATIONAL,
    feasibility_check,
    load_gnn_surrogate,
    traffic_light,
)
from pc_ddpm_alberta.modeling.predict import load_model, predict
from pc_ddpm_alberta.viz import REFERENCE_DAYS, build_overlay_figure, load_reference_day

DEFAULT_N_SCENARIOS = 24
MAX_N_SCENARIOS = 100

MODEL = load_model(device="cpu")
SURROGATE = load_gnn_surrogate(device="cpu")

_BADGE_COLOURS = {"green": "#10b981", "yellow": "#f59e0b", "red": "#ef4444"}


def _badge_html(report: dict[str, float | int]) -> str:
    colour = _BADGE_COLOURS[traffic_light(float(report["operational_pct"]))]
    op_lo, op_hi = V_BOUNDS_OPERATIONAL
    a_lo, a_hi = V_BOUNDS_ANSI
    return f"""
    <div style="border:1px solid #e5e7eb;border-radius:8px;padding:14px 16px;
                font-family:system-ui;line-height:1.5">
      <div style="display:flex;align-items:center;gap:10px;margin-bottom:6px">
        <span style="display:inline-block;width:14px;height:14px;border-radius:50%;
                     background:{colour}"></span>
        <strong>Voltage feasibility (GNN surrogate)</strong>
      </div>
      <div>Operational&nbsp;[{op_lo:.2f}, {op_hi:.2f}]&nbsp;pu:
        <strong>{report['operational_pct']:.1f}%</strong>
        ({report['operational_feasible']}/{report['n_total']})</div>
      <div>ANSI strict&nbsp;[{a_lo:.2f}, {a_hi:.2f}]&nbsp;pu:
        <strong>{report['ansi_pct']:.1f}%</strong>
        ({report['ansi_feasible']}/{report['n_total']})</div>
      <div style="color:#6b7280;font-size:12px;margin-top:6px">
        Fast proxy of the AC-PF headline metric; surrogate trained to
        &lt;3% MAPE on V_pu.
      </div>
    </div>
    """


def infer(
    reference_name: str,
    n_scenarios: int,
    progress: gr.Progress = gr.Progress(),
) -> tuple[matplotlib.figure.Figure, str]:
    reference = load_reference_day(reference_name)
    n_steps_total = int(MODEL.schedule["betas"].shape[0])

    def _cb(step: int, total: int) -> None:
        progress(step / total, desc=f"Reverse diffusion step {step}/{total}")

    progress(0.0, desc=f"Sampling {n_scenarios} scenarios over {n_steps_total} steps")
    scenarios = predict(MODEL, n_scenarios=int(n_scenarios), progress=_cb)

    progress(1.0, desc="Scoring feasibility with GNN surrogate")
    report = feasibility_check(scenarios, SURROGATE)

    return build_overlay_figure(scenarios, reference), _badge_html(report)


inputs = [
    gr.Radio(
        choices=list(REFERENCE_DAYS),
        value="typical_day",
        label="Reference day (overlay)",
        info="Real Alberta day plotted in red on top of the generated ensemble.",
    ),
    gr.Slider(
        minimum=4,
        maximum=MAX_N_SCENARIOS,
        value=DEFAULT_N_SCENARIOS,
        step=1,
        label="Number of scenarios",
        info="CPU latency on HF Spaces free tier: 4≈6s, 24≈38s, 100≈2.6min.",
    ),
]

outputs = [
    gr.Plot(label="Generated scenarios vs reference day"),
    gr.HTML(label="Voltage feasibility"),
]

examples = [
    ["typical_day", DEFAULT_N_SCENARIOS],
    ["high_wind_day", DEFAULT_N_SCENARIOS],
    ["low_wind_day", DEFAULT_N_SCENARIOS],
]


demo = gr.Interface(
    fn=infer,
    inputs=inputs,
    outputs=outputs,
    examples=examples,
    cache_examples=False,
    title="PC-DDPM Alberta — physics-constrained renewable energy scenario generation",
    description=(
        "Grid planners need scenarios that are statistically realistic AND physically "
        "consistent. PC-DDPM achieves 100% scenario feasibility under operational "
        "voltage bounds vs. 1.8% scenario infeasibility for the unconstrained baseline, "
        "on 9 years of real AESO data. "
        "Full methodology, code, and limitations: "
        "https://github.com/JBobyM/pc-ddpm-alberta · "
        "Model card: https://github.com/JBobyM/pc-ddpm-alberta/blob/master/model_card.md"
    ),
    article=(
        "See limitations and intended use in the "
        "[model card](https://github.com/JBobyM/pc-ddpm-alberta/blob/master/model_card.md)."
    ),
    flagging_mode="never",
)


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