Northern-Thailand-Ping-River-Monitor/src/ml/rain.py

"""Open-Meteo rainfall series for the upper Ping catchment.

One consistent source for training AND serving: the Open-Meteo forecast-model
archive (historical-forecast-api, 2021-03 onward) supplies hourly
precipitation at five catchment points above P.1; the live forecast endpoint
supplies the same series for recent days plus the next 48 h. Timestamps are
Asia/Bangkok local, matching the measurement grid. Rows before 2021-03 simply
have no rain data — HistGradientBoosting handles the NaNs natively.

The forward-looking sum built from this series is a legitimate *forecast*
feature, not label leakage: the series never contains river observations, and
at serving time the future values come from an actual weather forecast.
"""

import datetime
import logging
from pathlib import Path
from typing import Iterable, List, Optional, Tuple

import pandas as pd
import requests

logger = logging.getLogger(__name__)

# (name, lat, lon) — upper Ping catchment above P.1, headwaters to city
CATCHMENT_POINTS: Tuple[Tuple[str, float, float], ...] = (
    ("chiang_dao", 19.37, 98.97),
    ("mae_taeng", 19.12, 98.94),
    ("mae_ngat", 19.17, 99.05),
    ("mae_rim", 18.92, 98.92),
    ("chiang_mai", 18.79, 99.00),
)

HISTORY_URL = "https://historical-forecast-api.open-meteo.com/v1/forecast"
FORECAST_URL = "https://api.open-meteo.com/v1/forecast"
HISTORY_START = "2021-03-23"  # archive begins here
CACHE_FILE = "rain_openmeteo.csv.gz"


def _points_params() -> dict:
    return {
        "latitude": ",".join(str(lat) for _, lat, _ in CATCHMENT_POINTS),
        "longitude": ",".join(str(lon) for _, _, lon in CATCHMENT_POINTS),
        "hourly": "precipitation",
        "timezone": "Asia/Bangkok",
    }


def _parse_multi(payload, columns: Iterable[str]) -> pd.DataFrame:
    """Open-Meteo returns a list when multiple coordinates are requested."""
    results = payload if isinstance(payload, list) else [payload]
    frames = {}
    for name, result in zip(columns, results):
        hourly = result.get("hourly", {})
        idx = pd.to_datetime(hourly.get("time", []))
        frames[name] = pd.Series(hourly.get("precipitation", []), index=idx)
    df = pd.DataFrame(frames)
    df.index.name = "timestamp"
    return df


def fetch_history(
    start: str, end: str, session: Optional[requests.Session] = None
) -> pd.DataFrame:
    """Hourly precipitation for all catchment points over [start, end]."""
    session = session or requests.Session()
    response = session.get(
        HISTORY_URL,
        params={**_points_params(), "start_date": start, "end_date": end},
        timeout=120,
    )
    response.raise_for_status()
    return _parse_multi(response.json(), [p[0] for p in CATCHMENT_POINTS])


def fetch_forecast(
    past_days: int = 5,
    forecast_days: int = 2,
    session: Optional[requests.Session] = None,
) -> pd.DataFrame:
    """Recent + next-48h precipitation from the live forecast endpoint."""
    session = session or requests.Session()
    response = session.get(
        FORECAST_URL,
        params={
            **_points_params(),
            "past_days": past_days,
            "forecast_days": forecast_days,
        },
        timeout=60,
    )
    response.raise_for_status()
    return _parse_multi(response.json(), [p[0] for p in CATCHMENT_POINTS])


def load_history(
    cache_dir: Path = Path("models/cache"),
    end: Optional[datetime.date] = None,
    refresh: bool = True,
) -> Optional[pd.DataFrame]:
    """Cached catchment rain history from 2021-03 to ~today.

    Fetches year-sized chunks on first use (~6 requests), then only extends
    the tail. Returns None when the API is unreachable and no cache exists.
    """
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_dir / CACHE_FILE
    end = end or datetime.date.today()

    cached: Optional[pd.DataFrame] = None
    if cache_path.exists():
        cached = pd.read_csv(cache_path, index_col=0, parse_dates=True)

    fetch_from = pd.Timestamp(HISTORY_START)
    if cached is not None and len(cached):
        fetch_from = cached.index.max() - pd.Timedelta(days=2)  # re-fetch tail
    if not refresh and cached is not None:
        return cached

    chunks: List[pd.DataFrame] = []
    cursor = fetch_from.date()
    try:
        while cursor <= end:
            chunk_end = min(
                datetime.date(cursor.year, 12, 31), end
            )
            chunks.append(
                fetch_history(cursor.isoformat(), chunk_end.isoformat())
            )
            cursor = datetime.date(cursor.year + 1, 1, 1)
    except Exception as error:
        logger.warning(f"Open-Meteo history fetch failed: {error}")
        if not chunks and cached is None:
            return None

    if chunks:
        fresh = pd.concat(chunks)
        combined = (
            pd.concat([cached[cached.index < fresh.index.min()], fresh])
            if cached is not None
            else fresh
        )
        combined = combined[~combined.index.duplicated(keep="last")].sort_index()
        combined.to_csv(cache_path, compression="gzip")
        return combined
    return cached


def catchment_mean(df: Optional[pd.DataFrame]) -> Optional[pd.Series]:
    """Single catchment-average hourly rain series (mm)."""
    if df is None or df.empty:
        return None
    return df.mean(axis=1)


def serving_series() -> Optional[pd.Series]:
    """Catchment rain for inference: trailing days + the next 48 h forecast."""
    try:
        return catchment_mean(fetch_forecast())
    except Exception as error:
        logger.warning(f"Open-Meteo forecast fetch failed: {error}")
        return None


def backfill_db(engine, db_type: str, chunk_rows: int = 5000) -> int:
    """Push the full Open-Meteo history (2021+) into openmeteo_rain.

    Loads (or fetches) the archive cache and upserts in chunks; idempotent,
    safe to re-run, and safe alongside the hourly live writer.
    """
    history = load_history()
    if history is None or history.empty:
        logger.error("no rain history available to backfill")
        return 0
    total = 0
    for start in range(0, len(history), chunk_rows):
        part = history.iloc[start: start + chunk_rows]
        total += save_to_db(part, engine, db_type)
        logger.info(f"openmeteo_rain backfill: {total}/{len(history)} rows")
    return total


def save_to_db(df: pd.DataFrame, engine, db_type: str) -> int:
    """Upsert per-point + catchment-mean hourly rain into openmeteo_rain.

    Called by the leader worker's hourly precompute with the live forecast
    frame, so the DB accumulates both what fell (past rows are the model
    analysis) and what was forecast (future rows, overwritten as they become
    past). The ML training path reads Open-Meteo's own archive, not this
    table — this is for dashboards, SQL analysis, and source independence.
    """
    if df is None or df.empty:
        return 0
    from sqlalchemy import text

    point_cols = [p[0] for p in CATCHMENT_POINTS]
    ddl_cols = ", ".join(f"{c} NUMERIC(6,2)" for c in point_cols)
    ddl = (
        "CREATE TABLE IF NOT EXISTS openmeteo_rain ("
        "timestamp TIMESTAMP PRIMARY KEY, "
        f"{ddl_cols}, catchment_mean NUMERIC(6,2), "
        "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)"
    )
    cols = ["timestamp"] + point_cols + ["catchment_mean"]
    placeholders = ", ".join(f":{c}" for c in cols)
    updates = ", ".join(
        f"{c} = "
        + (f"VALUES({c})" if db_type == "mysql" else f"EXCLUDED.{c}")
        for c in cols[1:]
    )
    if db_type == "mysql":
        sql = (
            f"INSERT INTO openmeteo_rain ({', '.join(cols)}) VALUES ({placeholders}) "
            f"ON DUPLICATE KEY UPDATE {updates}"
        )
    else:
        sql = (
            f"INSERT INTO openmeteo_rain ({', '.join(cols)}) VALUES ({placeholders}) "
            f"ON CONFLICT (timestamp) DO UPDATE SET {updates}"
        )
    mean = df.mean(axis=1)
    params = [
        {
            "timestamp": ts.to_pydatetime(),
            **{c: (None if pd.isna(row[c]) else float(row[c])) for c in point_cols},
            "catchment_mean": None if pd.isna(mean.loc[ts]) else float(mean.loc[ts]),
        }
        for ts, row in df.iterrows()
    ]
    try:
        with engine.begin() as conn:
            conn.execute(text(ddl))
            conn.execute(text(sql), params)
        return len(params)
    except Exception as error:
        logger.error(f"openmeteo_rain save failed: {error}")
        return 0