weatherpred/dashboard.py

"""Gradio dashboard assembly."""

from __future__ import annotations

import gradio as gr
import pandas as pd
import plotly.graph_objects as go

from weatherpred.agent import make_gemini_brief, run_weather_agent
from weatherpred.modeling import build_weather_bundle
from weatherpred.zerogpu import spaces

def make_line_plot(
    comparison: pd.DataFrame,
    title: str,
    yaxis_title: str,
    series: list[tuple[str, str, str]],
) -> go.Figure:
    fig = go.Figure()
    for column, name, color in series:
        fig.add_trace(
            go.Scatter(
                x=comparison["date"],
                y=comparison[column],
                mode="lines+markers",
                name=name,
                line={"color": color, "width": 3},
            )
        )
    fig.update_layout(
        title=title,
        xaxis_title="Date",
        yaxis_title=yaxis_title,
        hovermode="x unified",
        template="plotly_white",
        legend={"orientation": "h", "y": -0.25},
        margin={"l": 40, "r": 20, "t": 60, "b": 70},
    )
    return fig


def make_temperature_plot(comparison: pd.DataFrame) -> go.Figure:
    return make_line_plot(
        comparison,
        "Five-Day Temperature Forecast",
        "Temperature (C)",
        [
            ("api_high_c", "Open-Meteo high", "#0f8b8d"),
            ("model_high_c", "ONNX model high", "#e05d3d"),
            ("api_low_c", "Open-Meteo low", "#74b6b7"),
            ("model_low_c", "ONNX model low", "#f0a08c"),
        ],
    )


def make_rain_plot(comparison: pd.DataFrame) -> go.Figure:
    return make_line_plot(
        comparison,
        "Five-Day Rain Forecast",
        "Rain (mm)",
        [
            ("api_rain_mm", "Open-Meteo rain", "#0f8b8d"),
            ("model_rain_mm", "Model rain proxy", "#e05d3d"),
        ],
    )


def make_wind_plot(comparison: pd.DataFrame) -> go.Figure:
    return make_line_plot(
        comparison,
        "Five-Day Wind Forecast",
        "Wind (km/h)",
        [
            ("api_wind_kmh", "Open-Meteo wind", "#0f8b8d"),
            ("model_wind_kmh", "Historical model wind", "#e05d3d"),
        ],
    )


@spaces.GPU(duration=10)
def zero_gpu_probe() -> str:
    return "ZeroGPU compatibility probe is available."


def make_metrics(summary: dict) -> str:
    return (
        f"Historical rows: {summary['history_rows']:,}\n"
        f"History window: {summary['history_start']} to {summary['history_end']}\n"
        f"Average model confidence proxy: {summary['avg_model_confidence']:.0%}\n"
        f"Open-Meteo max high: {summary['max_api_high_c']:.1f}C\n"
        f"ONNX model max high: {summary['max_model_high_c']:.1f}C"
    )


def render_all_charts(comparison: pd.DataFrame) -> tuple:
    return (
        gr.update(value=make_temperature_plot(comparison), visible=True),
        gr.update(value=make_rain_plot(comparison), visible=True),
        gr.update(value=make_wind_plot(comparison), visible=True),
    )


def refresh_dashboard(force_refresh_history: bool) -> tuple:
    bundle = build_weather_bundle(force_refresh_history=force_refresh_history)
    comparison = bundle.comparison.copy()
    summary = bundle.summary
    brief = make_gemini_brief(comparison, summary)
    return (
        comparison,
        "all",
        make_temperature_plot(comparison),
        make_rain_plot(comparison),
        make_wind_plot(comparison),
        make_metrics(summary),
        brief,
        [{"role": "assistant", "content": "Dashboard refreshed. Showing temperature, rain, and wind forecasts."}],
    )


def parse_chart_selection(message: str) -> tuple[str, str]:
    text = message.lower()
    wants_temperature = any(word in text for word in ["temperature", "temp", "hot", "cold", "high", "low"])
    wants_rain = any(word in text for word in ["rain", "precipitation", "precip", "shower"])
    wants_wind = "wind" in text
    wants_all = any(word in text for word in ["all", "everything", "reset", "both"])

    if wants_all or sum([wants_temperature, wants_rain, wants_wind]) > 1:
        return "all", "Showing all three charts in the fixed 2x2 layout."
    if wants_temperature:
        return "temperature", "Temperature is shown in the top-left chart."
    if wants_rain:
        return "rain", "Rain is shown in the top-right chart."
    if wants_wind:
        return "wind", "Wind is shown in the bottom-left chart, including the historical model wind prediction."
    return "all", "I can point you to temperature, rain, wind, or all charts in the 2x2 dashboard."


def chat_with_dashboard(message: str, history: list, comparison: pd.DataFrame) -> tuple:
    if comparison is None or getattr(comparison, "empty", True):
        comparison = build_weather_bundle(force_refresh_history=False).comparison.copy()

    if not message.strip():
        return (history or [], "all", comparison, *render_all_charts(comparison), "")

    agent_run = run_weather_agent(message)
    next_comparison = agent_run.comparison if agent_run.comparison is not None else comparison
    charts = render_all_charts(next_comparison)
    history = history or []
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": agent_run.response})
    selection = agent_run.action.get("variable", "all")
    return (history, selection, next_comparison, *charts, agent_run.response)


with gr.Blocks(title="Tokyo Weather Forecast Lab") as demo:
    gr.Markdown(
        """
        # Tokyo Weather Forecast Lab

        Compare Open-Meteo's five-day forecast with a local ONNX model using
        cached 30-year Tokyo weather history. Use the chat panel to show
        temperature, rain, wind, or all forecast charts in the fixed 2x2 layout.
        """
    )
    forecast_state = gr.State(pd.DataFrame())
    chart_selection = gr.State("all")

    with gr.Row():
        refresh = gr.Button("Refresh forecast", variant="primary")
        force_history = gr.Checkbox(
            label="Refresh 30-year historical cache",
            value=False,
        )

    with gr.Row(visible=False):
        gpu_probe_button = gr.Button("ZeroGPU probe")
        gpu_probe_output = gr.Textbox(label="ZeroGPU probe")

    with gr.Row():
        with gr.Column(scale=1):
            temperature_plot = gr.Plot(label="Temperature")
        with gr.Column(scale=1):
            rain_plot = gr.Plot(label="Rain")

    with gr.Row():
        with gr.Column(scale=1):
            wind_plot = gr.Plot(label="Wind")
        with gr.Column(scale=1):
            gr.Markdown("## Forecast Chat")
            chatbot = gr.Chatbot(label="Dashboard assistant", height=360)
            chat_input = gr.Textbox(
                label="Ask the dashboard",
                placeholder="Try: show rain, show wind, show temperature, show all charts",
            )

    with gr.Row():
        metrics_box = gr.Textbox(label="Data and model metrics", lines=6)
        brief_box = gr.Textbox(label="AI analyst brief", lines=6)

    gpu_probe_button.click(
        zero_gpu_probe,
        inputs=None,
        outputs=gpu_probe_output,
    )

    refresh.click(
        refresh_dashboard,
        inputs=[force_history],
        outputs=[
            forecast_state,
            chart_selection,
            temperature_plot,
            rain_plot,
            wind_plot,
            metrics_box,
            brief_box,
            chatbot,
        ],
    )
    demo.load(
        refresh_dashboard,
        inputs=[force_history],
        outputs=[
            forecast_state,
            chart_selection,
            temperature_plot,
            rain_plot,
            wind_plot,
            metrics_box,
            brief_box,
            chatbot,
        ],
    )
    chat_input.submit(
        chat_with_dashboard,
        inputs=[chat_input, chatbot, forecast_state],
        outputs=[
            chatbot,
            chart_selection,
            forecast_state,
            temperature_plot,
            rain_plot,
            wind_plot,
            brief_box,
        ],
    ).then(lambda: "", outputs=chat_input)