weatherpred/tools.py

"""Deterministic tools called by the weather dashboard agent."""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

from weatherpred.config import FORECAST_DAYS
from weatherpred.modeling import build_weather_bundle

MAX_FORECAST_DAYS = 16
VARIABLE_COLUMNS = {
    "temperature": {
        "api": ["api_high_c", "api_low_c"],
        "model": ["model_high_c", "model_low_c"],
        "history": "high_c",
        "unit": "C",
    },
    "rain": {
        "api": ["api_rain_mm"],
        "model": ["model_rain_mm"],
        "history": "rain_mm",
        "unit": "mm",
    },
    "wind": {
        "api": ["api_wind_kmh"],
        "model": ["model_wind_kmh"],
        "history": "wind_kmh",
        "unit": "km/h",
    },
}


def normalize_variable(variable: str | None) -> str:
    text = (variable or "all").lower()
    if text in {"temp", "temperature", "hot", "cold"}:
        return "temperature"
    if text in {"rain", "precip", "precipitation", "shower"}:
        return "rain"
    if text in {"wind", "windspeed", "wind_speed"}:
        return "wind"
    return "all"


def clamp_horizon(horizon_days: int | str | None) -> int:
    try:
        value = int(horizon_days or FORECAST_DAYS)
    except (TypeError, ValueError):
        value = FORECAST_DAYS
    return max(1, min(MAX_FORECAST_DAYS, value))


def compact_forecast_rows(comparison: pd.DataFrame, variable: str) -> list[dict[str, Any]]:
    variable = normalize_variable(variable)
    columns = ["date"]
    selected = ["temperature", "rain", "wind"] if variable == "all" else [variable]
    for item in selected:
        columns.extend(VARIABLE_COLUMNS[item]["api"])
        columns.extend(VARIABLE_COLUMNS[item]["model"])
    return comparison[columns].round(2).to_dict(orient="records")


def forecast_next_days(variable: str = "all", horizon_days: int = FORECAST_DAYS) -> dict[str, Any]:
    variable = normalize_variable(variable)
    horizon = clamp_horizon(horizon_days)
    bundle = build_weather_bundle(horizon_days=horizon)
    return {
        "tool": "forecast_next_days",
        "variable": variable,
        "horizon_days": horizon,
        "comparison": bundle.comparison,
        "summary": {
            **bundle.summary,
            "rows": compact_forecast_rows(bundle.comparison, variable),
            "model_note": "ONNX predicts max temperature for the first five days; other outputs use historical seasonal statistics.",
        },
    }


def analyze_historical_trend(variable: str = "temperature") -> dict[str, Any]:
    from weatherpred.data import load_or_fetch_history

    variable = normalize_variable(variable)
    if variable == "all":
        variable = "temperature"
    info = VARIABLE_COLUMNS[variable]
    history = load_or_fetch_history(force_refresh=False).sort_values("date")
    yearly = history.set_index("date")[info["history"]].resample("YE").mean().dropna()
    x = np.arange(len(yearly), dtype=float)
    slope = float(np.polyfit(x, yearly.to_numpy(dtype=float), 1)[0]) if len(yearly) > 1 else 0.0
    first_mean = float(yearly.head(5).mean())
    last_mean = float(yearly.tail(5).mean())
    return {
        "tool": "analyze_historical_trend",
        "variable": variable,
        "summary": {
            "history_start": history["date"].min().strftime("%Y-%m-%d"),
            "history_end": history["date"].max().strftime("%Y-%m-%d"),
            "unit": info["unit"],
            "first_5_year_mean": round(first_mean, 2),
            "last_5_year_mean": round(last_mean, 2),
            "change_last_vs_first": round(last_mean - first_mean, 2),
            "linear_slope_per_year": round(slope, 3),
        },
    }


def analyze_seasonality(variable: str = "temperature") -> dict[str, Any]:
    from weatherpred.data import load_or_fetch_history

    variable = normalize_variable(variable)
    if variable == "all":
        variable = "temperature"
    info = VARIABLE_COLUMNS[variable]
    history = load_or_fetch_history(force_refresh=False).copy()
    history["month"] = history["date"].dt.month
    monthly = history.groupby("month")[info["history"]].mean()
    peak_month = int(monthly.idxmax())
    low_month = int(monthly.idxmin())
    return {
        "tool": "analyze_seasonality",
        "variable": variable,
        "summary": {
            "unit": info["unit"],
            "peak_month": peak_month,
            "peak_value": round(float(monthly.loc[peak_month]), 2),
            "low_month": low_month,
            "low_value": round(float(monthly.loc[low_month]), 2),
            "seasonal_amplitude": round(float(monthly.max() - monthly.min()), 2),
            "monthly_means": {str(month): round(float(value), 2) for month, value in monthly.items()},
        },
    }


def compare_model_vs_open_meteo(variable: str = "all", horizon_days: int = FORECAST_DAYS) -> dict[str, Any]:
    forecast = forecast_next_days(variable=variable, horizon_days=horizon_days)
    comparison = forecast["comparison"]
    variable = normalize_variable(variable)
    selected = ["temperature", "rain", "wind"] if variable == "all" else [variable]
    deltas: dict[str, Any] = {}
    for item in selected:
        if item == "temperature":
            deltas["temperature_high_mae_c"] = round(float((comparison["model_high_c"] - comparison["api_high_c"]).abs().mean()), 2)
            deltas["temperature_low_mae_c"] = round(float((comparison["model_low_c"] - comparison["api_low_c"]).abs().mean()), 2)
        elif item == "rain":
            deltas["rain_mae_mm"] = round(float((comparison["model_rain_mm"] - comparison["api_rain_mm"]).abs().mean()), 2)
        elif item == "wind":
            deltas["wind_mae_kmh"] = round(float((comparison["model_wind_kmh"] - comparison["api_wind_kmh"]).abs().mean()), 2)
    forecast["tool"] = "compare_model_vs_open_meteo"
    forecast["summary"]["average_absolute_differences"] = deltas
    return forecast


def explain_model() -> dict[str, Any]:
    return {
        "tool": "explain_model",
        "variable": "all",
        "summary": {
            "data_source": "Open-Meteo archive API for history and Open-Meteo forecast API as an external baseline.",
            "history_window": "Up to 30 years of daily Tokyo max/min temperature, precipitation, and max wind speed stored in local CSV/JSON cache.",
            "model": "A small ONNX demo model loaded from models/mock_model/model.onnx.",
            "direct_prediction": "The ONNX model directly predicts five days of max temperature from the previous seven daily max temperatures.",
            "derived_outputs": "Min temperature, rain, wind, confidence, and horizons beyond five days are derived from historical seasonal statistics.",
            "limitation": "This is an interview/demo forecasting scaffold, not a production meteorological model.",
        },
    }


def show_dashboard_view(variable: str = "all") -> dict[str, Any]:
    variable = normalize_variable(variable)
    locations = {
        "temperature": "top-left chart",
        "rain": "top-right chart",
        "wind": "bottom-left chart",
        "all": "the three chart cells in the 2x2 layout",
    }
    return {
        "tool": "show_dashboard_view",
        "variable": variable,
        "summary": {"view": locations[variable]},
    }


def run_tool(action: dict[str, Any]) -> dict[str, Any]:
    tool = action.get("tool", "show_dashboard_view")
    variable = normalize_variable(action.get("variable"))
    horizon = clamp_horizon(action.get("horizon_days"))
    if tool == "forecast_next_days":
        return forecast_next_days(variable, horizon)
    if tool == "analyze_historical_trend":
        return analyze_historical_trend(variable)
    if tool == "analyze_seasonality":
        return analyze_seasonality(variable)
    if tool == "compare_model_vs_open_meteo":
        return compare_model_vs_open_meteo(variable, horizon)
    if tool == "explain_model":
        return explain_model()
    return show_dashboard_view(variable)