timesfm_wrapper.py

"""
timesfm_wrapper.py
=============================
"""

from __future__ import annotations

import hashlib
import logging
import math
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional, Protocol, Tuple, runtime_checkable

import numpy as np

import zone_observation as _zo

assert _zo.SCHEMA_VERSION == 3, (
    f"timesfm_wrapper: zone_observation schema mismatch "
    f"(expected 3, got {_zo.SCHEMA_VERSION})"
)

from zone_observation import (
    DataSource,
    ForecastResult,
    ZoneObs,
    _stable_seed,
)

logger = logging.getLogger(__name__)

try:
    import requests as _requests
    REQUESTS_AVAILABLE = True
except ImportError:
    REQUESTS_AVAILABLE = False
    logger.info("requests not installed — OpenMeteoBackend unavailable")


# ---------------------------------------------------------------------------
# Protocol
# ---------------------------------------------------------------------------

@runtime_checkable
class ForecastBackend(Protocol):
    def forecast(self, obs: ZoneObs) -> ForecastResult:
        ...


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _clip_prob(values: List[float]) -> Tuple[float, ...]:
    return tuple(max(0.0, min(1.0, v)) for v in values)


def _make_temp_sequences(
    temp_mean_c: float,
    horizon_days: int,
    spread_c: float = 2.0,
) -> Tuple[Tuple[float, ...], Tuple[float, ...], Tuple[float, ...]]:
    mean_seq = tuple(float(temp_mean_c)             for _ in range(horizon_days))
    p10_seq  = tuple(float(temp_mean_c - spread_c)  for _ in range(horizon_days))
    p90_seq  = tuple(float(temp_mean_c + spread_c)  for _ in range(horizon_days))
    return mean_seq, p10_seq, p90_seq


def _rh_to_prob_humidity(rh_max_pct: float) -> float:
    x = (rh_max_pct - 80.0) / 10.0
    return float(max(0.0, min(1.0, 1.0 / (1.0 + math.exp(-x)))))


# ---------------------------------------------------------------------------
# Tier 1 — BaselineBackend
# ---------------------------------------------------------------------------

class BaselineBackend:
    """
    Deterministic statistical forecast derived entirely from ZoneObs fields.

    No randomness — deterministic forecasts let the policy learn stable
    input-output mappings. Stochasticity enters through the observation
    pipeline, not the forecast.
    """

    def __init__(self, horizon_days: int = 30) -> None:
        self.horizon_days = horizon_days

    def forecast(self, obs: ZoneObs) -> ForecastResult:
        h = self.horizon_days

        daily_avg = obs.precip_30d_mm / max(30.0, 1.0)
        decay     = np.linspace(1.0, 0.3, h)
        precip    = np.clip(daily_avg * decay, 0.0, 500.0)

        flood_m   = 1.0 + 0.8 * obs.flood_signal()
        drought_m = 1.0 - 0.6 * obs.drought_signal()
        precip    = np.clip(precip * flood_m * drought_m, 0.0, 500.0)

        lo_factor = np.linspace(0.75, 0.50, h)
        hi_factor = np.linspace(1.25, 1.50, h)
        p10 = np.clip(precip * lo_factor, 0.0, 500.0)
        p90 = np.clip(precip * hi_factor, 0.0, 500.0)

        prob_heavy   = [float(min(1.0, v / 20.0))       for v in precip]
        prob_drought = [float(max(0.0, 1.0 - v / 2.0))  for v in precip]
        prob_humid   = [_rh_to_prob_humidity(obs.rh_max_pct)] * h

        temp_mean, temp_p10, temp_p90 = _make_temp_sequences(obs.temp_mean_c, h)
        rh_seq = tuple(float(obs.rh_mean_pct) for _ in range(h))

        return ForecastResult(
            zone_id=obs.zone_id,
            forecast_time=obs.valid_time,
            horizon_days=h,
            precip_mm=tuple(float(v) for v in precip),
            precip_p10=tuple(float(v) for v in p10),
            precip_p90=tuple(float(v) for v in p90),
            temp_mean_c=temp_mean,
            temp_p10=temp_p10,
            temp_p90=temp_p90,
            rh_mean_pct=rh_seq,
            prob_heavy_rain=_clip_prob(prob_heavy),
            prob_drought_day=_clip_prob(prob_drought),
            prob_high_humidity=_clip_prob(prob_humid),
            model_id="baseline-persistence-v1",
            source=DataSource.SYNTHETIC,
        )


# ---------------------------------------------------------------------------
# Tier 2 — OpenMeteoBackend
# ---------------------------------------------------------------------------

class OpenMeteoBackend:
    """
    Real probabilistic forecast via Open-Meteo free API.

    FIX vs original: lat/lon moved from forecast() into __init__() so the
    object satisfies the ForecastBackend Protocol (forecast(self, obs) only).
    Resolve lat/lon from the zone registry before constructing this object.

    Example:
        lat, lon = zone_registry[zone_id].centroid
        backend  = OpenMeteoBackend(lat=lat, lon=lon, horizon_days=16)
        result   = backend.forecast(obs)
    """

    BASE_URL    = "https://api.open-meteo.com/v1/forecast"
    TIMEOUT_S   = 15

    def __init__(self, lat: float, lon: float, horizon_days: int = 16) -> None:
        self.lat          = lat
        self.lon          = lon
        self.horizon_days = min(horizon_days, 16)
        if horizon_days > 16:
            logger.warning(
                "OpenMeteoBackend: horizon_days=%d capped at 16 "
                "(Open-Meteo free tier limit).", horizon_days
            )

    def forecast(self, obs: ZoneObs) -> ForecastResult:
        if not REQUESTS_AVAILABLE:
            raise RuntimeError(
                "requests not installed — cannot use OpenMeteoBackend. "
                "Install with: pip install requests"
            )

        h = self.horizon_days
        params = {
            "latitude":  self.lat,
            "longitude": self.lon,
            "daily": ",".join([
                "precipitation_sum",
                "temperature_2m_mean",
                "temperature_2m_max",
                "temperature_2m_min",
                "relative_humidity_2m_mean",
                "relative_humidity_2m_max",
                "precipitation_probability_mean",
                "et0_fao_evapotranspiration",
            ]),
            "forecast_days": h,
            "timezone":      "UTC",
        }

        try:
            resp = _requests.get(self.BASE_URL, params=params, timeout=self.TIMEOUT_S)
            resp.raise_for_status()
            data = resp.json()
        except Exception as e:
            logger.warning(
                "OpenMeteoBackend: API call failed (%s) — falling back to BaselineBackend", e
            )
            return BaselineBackend(horizon_days=h).forecast(obs)

        daily = data.get("daily", {})

        def _safe(key: str, n: int, default: float = 0.0) -> List[float]:
            vals = daily.get(key, [])
            return [
                float(vals[i]) if i < len(vals) and vals[i] is not None else default
                for i in range(n)
            ]

        precip_raw  = _safe("precipitation_sum",              h)
        temp_mean   = _safe("temperature_2m_mean",            h, obs.temp_mean_c)
        temp_max    = _safe("temperature_2m_max",             h, obs.temp_max_c)
        temp_min    = _safe("temperature_2m_min",             h, obs.temp_min_c)
        rh_mean     = _safe("relative_humidity_2m_mean",      h, obs.rh_mean_pct)
        rh_max_raw  = _safe("relative_humidity_2m_max",       h, obs.rh_max_pct)
        precip_prob = _safe("precipitation_probability_mean", h, 50.0)

        precip = [max(0.0, min(500.0, v)) for v in precip_raw]
        t_p10  = [min(mn, mx) for mn, mx in zip(temp_min, temp_mean)]
        t_p90  = [max(mx, mn) for mx, mn in zip(temp_max, temp_mean)]

        obs_spread = float(np.clip(
            (obs.precip_30d_mm - obs.precip_7d_mm * 4.0)
            / max(obs.precip_30d_mm, 1.0),
            0.1, 0.6,
        ))
        unc = [obs_spread * (1.0 + 0.5 * i / h) for i in range(h)]
        p10 = [max(0.0,   p * (1.0 - u * 0.8)) for p, u in zip(precip, unc)]
        p90 = [min(500.0, p * (1.0 + u))        for p, u in zip(precip, unc)]

        prob_heavy   = [max(0.0, min(1.0, pp / 100.0))       for pp in precip_prob]
        prob_drought = [max(0.0, min(1.0, 1.0 - pp / 100.0)) for pp in precip_prob]
        prob_humid   = [_rh_to_prob_humidity(rh)              for rh in rh_max_raw]

        return ForecastResult(
            zone_id=obs.zone_id,
            forecast_time=obs.valid_time,
            horizon_days=h,
            precip_mm=tuple(precip),
            precip_p10=tuple(p10),
            precip_p90=tuple(p90),
            temp_mean_c=tuple(temp_mean),
            temp_p10=tuple(t_p10),
            temp_p90=tuple(t_p90),
            rh_mean_pct=tuple(rh_mean),
            prob_heavy_rain=_clip_prob(prob_heavy),
            prob_drought_day=_clip_prob(prob_drought),
            prob_high_humidity=_clip_prob(prob_humid),
            model_id="open-meteo-forecast-v1",
            source=DataSource.OPENMETEO_LIVE,
        )


# ---------------------------------------------------------------------------
# Tier 3 — LocalTimesFMBackend
# ---------------------------------------------------------------------------

@dataclass
class LocalTimesFMBackend:
    """
    Local TimesFM-2.5-200m inference. SHA256-verified. No synthetic fallback.

    Download checkpoint from:
      https://huggingface.co/google/timesfm-2.5-200m-pytorch
    """

    checkpoint_path: Path
    expected_sha256: str
    device:          str = "cpu"
    horizon_days:    int = 30
    context_len:     int = 512

    _model: Optional[object] = field(default=None, init=False, repr=False)

    def __post_init__(self) -> None:
        self.checkpoint_path = Path(self.checkpoint_path)
        if not self.checkpoint_path.exists():
            raise FileNotFoundError(
                f"TimesFM checkpoint not found: {self.checkpoint_path}"
            )
        logger.info("Verifying TimesFM checkpoint SHA256...")
        computed = self._sha256_file(self.checkpoint_path)
        if computed != self.expected_sha256.lower().strip():
            raise RuntimeError(
                f"TimesFM checkpoint SHA256 MISMATCH.\n"
                f"Expected: {self.expected_sha256}\nGot: {computed}"
            )
        logger.info("TimesFM checkpoint verified OK")
        self._model = self._load_model()

    def _load_model(self) -> object:
        try:
            import timesfm  # type: ignore
        except ImportError:
            raise ImportError("timesfm not installed. pip install timesfm")
        model = timesfm.TimesFm(
            hparams=timesfm.TimesFmHparams(
                backend=self.device,
                per_core_batch_size=32,
                horizon_len=self.horizon_days,
                context_len=self.context_len,
            ),
            checkpoint=timesfm.TimesFmCheckpoint(
                local_path=str(self.checkpoint_path)
            ),
        )
        logger.info("TimesFM loaded: device=%s  horizon=%d", self.device, self.horizon_days)
        return model

    def forecast(self, obs: ZoneObs) -> ForecastResult:
        if self._model is None:
            raise RuntimeError("TimesFM model not loaded.")

        import pandas as pd  # type: ignore
        import datetime as _dt

        h = self.horizon_days

        daily_30d = obs.precip_30d_mm / 30.0
        daily_7d  = obs.precip_7d_mm  / 7.0
        daily_24h = obs.precip_24h_mm

        context = np.concatenate([
            np.full(23, daily_30d),
            np.full(6,  daily_7d),
            np.array([daily_24h]),
        ]).astype(np.float32)

        base_date = obs.valid_time.date()
        dates = [
            base_date - _dt.timedelta(days=len(context) - i - 1)
            for i in range(len(context))
        ]
        df_input = pd.DataFrame({
            "unique_id": [obs.zone_id] * len(context),
            "ds":        dates,
            "y":         context,
        })

        forecast_df, _ = self._model.forecast_on_df(
            inputs=df_input,
            freq="D",
            value_name="y",
            num_jobs=1,
        )

        p50_col = "timesfm-q-0.5" if "timesfm-q-0.5" in forecast_df.columns else "timesfm"
        p10_col = "timesfm-q-0.1" if "timesfm-q-0.1" in forecast_df.columns else p50_col
        p90_col = "timesfm-q-0.9" if "timesfm-q-0.9" in forecast_df.columns else p50_col

        p50 = np.clip(forecast_df[p50_col].values[:h].astype(float), 0.0, 500.0)
        p10 = np.clip(forecast_df[p10_col].values[:h].astype(float), 0.0, 500.0)
        p90 = np.clip(forecast_df[p90_col].values[:h].astype(float), 0.0, 500.0)

        p10 = np.minimum(p10, p50)
        p90 = np.maximum(p90, p50)

        prob_heavy   = [float(min(1.0, v / 20.0))      for v in p50]
        prob_drought = [float(max(0.0, 1.0 - v / 2.0)) for v in p50]
        prob_humid   = [_rh_to_prob_humidity(obs.rh_max_pct)] * h

        temp_mean_seq, temp_p10_seq, temp_p90_seq = _make_temp_sequences(
            obs.temp_mean_c, h
        )

        return ForecastResult(
            zone_id=obs.zone_id,
            forecast_time=obs.valid_time,
            horizon_days=h,
            precip_mm=tuple(float(v) for v in p50),
            precip_p10=tuple(float(v) for v in p10),
            precip_p90=tuple(float(v) for v in p90),
            temp_mean_c=temp_mean_seq,
            temp_p10=temp_p10_seq,
            temp_p90=temp_p90_seq,
            rh_mean_pct=tuple(float(obs.rh_mean_pct) for _ in range(h)),
            prob_heavy_rain=_clip_prob(prob_heavy),
            prob_drought_day=_clip_prob(prob_drought),
            prob_high_humidity=_clip_prob(prob_humid),
            model_id="timesfm-2.5-200m",
            source=DataSource.SYNTHETIC,   # FIX: was ERA5_REANALYSIS (incorrect)
        )

    @staticmethod
    def _sha256_file(path: Path, chunk_size: int = 1 << 20) -> str:
        h = hashlib.sha256()
        with open(path, "rb") as f:
            while chunk := f.read(chunk_size):
                h.update(chunk)
        return h.hexdigest()


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

def create_forecast_backend(
    mode:            str            = "baseline",
    checkpoint_path: Optional[Path] = None,
    expected_sha256: Optional[str]  = None,
    device:          str            = "cpu",
    horizon_days:    int            = 30,
    lat:             Optional[float] = None,
    lon:             Optional[float] = None,
) -> ForecastBackend:
    """
    Create a ForecastBackend.

    mode="baseline"   — deterministic statistical forecast (default, no deps)
    mode="openmeteo"  — real Open-Meteo API forecast (requires requests, lat, lon)
    mode="timesfm"    — local TimesFM inference (requires checkpoint + timesfm)
    """
    if mode == "baseline":
        return BaselineBackend(horizon_days=horizon_days)

    if mode == "openmeteo":
        if not REQUESTS_AVAILABLE:
            logger.warning(
                "requests not installed — falling back to BaselineBackend."
            )
            return BaselineBackend(horizon_days=horizon_days)
        if lat is None or lon is None:
            raise ValueError(
                "mode='openmeteo' requires lat and lon. "
                "Resolve from zone registry: lat, lon = registry[zone_id].centroid"
            )
        return OpenMeteoBackend(lat=lat, lon=lon, horizon_days=horizon_days)

    if mode == "timesfm":
        if checkpoint_path is None or expected_sha256 is None:
            raise ValueError(
                "mode='timesfm' requires both checkpoint_path and expected_sha256."
            )
        return LocalTimesFMBackend(
            checkpoint_path=Path(checkpoint_path),
            expected_sha256=expected_sha256,
            device=device,
            horizon_days=horizon_days,
        )

    raise ValueError(
        f"Unknown backend mode: '{mode}'. Valid: 'baseline', 'openmeteo', 'timesfm'"
    )