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

"""Mae Ngat reservoir series for the flood models.

rid_reservoir_daily (collected hourly by src/rid_reservoir.py, backfilled to
2018) holds daily storage/inflow/outflow for every RID large dam. Mae Ngat
Somboon Chon (DAM_ID 200103) is the only large dam upstream of Chiang Mai:
in Oct 2024 its inflow hit 19-22 MCM/day and storage 114% of usable capacity
days around the P.1 crossing — upstream state no river gauge carries.

Leakage rule: RID publishes the daily report for date D on the morning of D,
so the row becomes visible to features at D 07:00 local time, never earlier.
Forward-fill is capped at FFILL_LIMIT_H so a stalled collector degrades to
NaN (HGB-native) instead of silently serving stale reservoir state.

Known residual optimism: the collector upserts keep-last (and re-fetches
yesterday), so the stored row for date D is RID's FINAL revision, which
training then back-dates to D 07:00 — values live serving may not have had
that morning. This bias works IN FAVOR of dam features, so the 2026-08-13
negative result (they cost 1-3 h of alert lead) holds a fortiori; but any
future POSITIVE result must first validate intraday row stability or shift
the flow columns to D+1 07:00.
"""

import datetime
import logging
from pathlib import Path
from typing import Optional

import pandas as pd

from ..rid_reservoir import MAE_NGAT_DAM_ID
from .data import CACHE_DIR, resolve_db_url

logger = logging.getLogger(__name__)

REPORT_HOUR = 7  # daily value valid from 07:00 local on its own date
FFILL_LIMIT_H = 48  # two missed daily reports -> NaN, not stale state
DAM_COLUMNS = ("storage_pct", "inflow_mcm", "outflow_mcm")
CACHE_FILE = f"dam_{MAE_NGAT_DAM_ID}.csv.gz"


def load_daily(
    db_url: Optional[str] = None,
    dam_id: str = MAE_NGAT_DAM_ID,
    start: Optional[datetime.date] = None,
    cache_dir: Path = CACHE_DIR,
) -> Optional[pd.DataFrame]:
    """Daily dam rows indexed by date. DB first, on-disk cache as fallback."""
    cache_path = Path(cache_dir) / CACHE_FILE
    resolved = resolve_db_url(db_url)
    if resolved:
        try:
            from sqlalchemy import create_engine, text

            query = (
                "SELECT date, storage_pct, inflow_mcm, outflow_mcm "
                "FROM rid_reservoir_daily WHERE dam_id = :dam_id"
            )
            params = {"dam_id": dam_id}
            if start is not None:
                query += " AND date >= :start"
                params["start"] = start
            engine = create_engine(resolved, pool_pre_ping=True)
            with engine.connect() as conn:
                daily = pd.read_sql(
                    text(query + " ORDER BY date"), conn, params=params
                )
            daily["date"] = pd.to_datetime(daily["date"])
            daily = daily.set_index("date")
            for col in DAM_COLUMNS:
                daily[col] = pd.to_numeric(daily[col], errors="coerce")
            # Only full, NON-EMPTY loads refresh the cache: a truncated or
            # freshly-recreated table must not wipe a good fallback archive.
            if start is None and not daily.empty:
                cache_path.parent.mkdir(parents=True, exist_ok=True)
                daily.to_csv(cache_path, compression="gzip")
            return daily
        except Exception as error:
            logger.warning(f"dam series DB load failed: {error}")
    if cache_path.exists():
        logger.warning("falling back to on-disk cache for the dam series")
        return pd.read_csv(cache_path, index_col=0, parse_dates=True)
    return None


def hourly_frame(daily: Optional[pd.DataFrame]) -> Optional[pd.DataFrame]:
    """Step the daily rows onto an hourly grid, each valid from D 07:00."""
    if daily is None or daily.empty:
        return None
    frame = daily.copy()
    frame.index = pd.to_datetime(frame.index) + pd.Timedelta(hours=REPORT_HOUR)
    frame = frame[~frame.index.duplicated(keep="last")].sort_index()
    hourly_index = pd.date_range(
        frame.index.min(),
        frame.index.max() + pd.Timedelta(hours=FFILL_LIMIT_H),
        freq="h",
    )
    return frame.reindex(hourly_index).ffill(limit=FFILL_LIMIT_H)


def load_history(db_url: Optional[str] = None) -> Optional[pd.DataFrame]:
    """Full hourly Mae Ngat history for training; None when unavailable."""
    return hourly_frame(load_daily(db_url))


def serving_frame(
    db_url: Optional[str] = None, days: int = 21
) -> Optional[pd.DataFrame]:
    """Recent hourly dam state for inference (covers the 336 h feature window
    plus the 72 h storage-delta lag)."""
    start = datetime.date.today() - datetime.timedelta(days=days)
    return hourly_frame(load_daily(db_url, start=start))