Northern-Thailand-Ping-River-Monitor/tests/test_dam_features.py

"""Tests for the Mae Ngat dam series (src/ml/dam.py) and its feature gating."""

import datetime

import numpy as np
import pandas as pd

from src.ml import features
from src.ml.dam import FFILL_LIMIT_H, REPORT_HOUR, hourly_frame


def _daily(days=5, start="2024-09-20"):
    idx = pd.date_range(start, periods=days, freq="D")
    return pd.DataFrame(
        {
            "storage_pct": np.linspace(90, 110, days),
            "inflow_mcm": np.linspace(2, 20, days),
            "outflow_mcm": np.linspace(0.5, 5, days),
        },
        index=idx,
    )


class TestHourlyFrame:
    def test_daily_value_visible_from_report_hour_only(self):
        hourly = hourly_frame(_daily())
        day0 = pd.Timestamp("2024-09-20")
        # Nothing before the first report hour
        assert hourly.index.min() == day0 + pd.Timedelta(hours=REPORT_HOUR)
        # The day's value holds from 07:00 through the next morning
        assert hourly.loc[day0 + pd.Timedelta(hours=7), "storage_pct"] == 90.0
        assert hourly.loc[day0 + pd.Timedelta(hours=23), "storage_pct"] == 90.0
        next_6am = day0 + pd.Timedelta(days=1, hours=6)
        next_7am = day0 + pd.Timedelta(days=1, hours=7)
        assert hourly.loc[next_6am, "storage_pct"] == 90.0  # yesterday's value
        assert hourly.loc[next_7am, "storage_pct"] == 95.0  # today's report

    def test_ffill_capped_after_missing_days(self):
        daily = _daily(days=2).drop(index=pd.Timestamp("2024-09-21"))
        # extend with a far-later row so the gap sits mid-frame
        late = _daily(days=1, start="2024-09-28")
        hourly = hourly_frame(pd.concat([daily, late]))
        gap_ts = pd.Timestamp("2024-09-20") + pd.Timedelta(
            hours=REPORT_HOUR + FFILL_LIMIT_H + 1
        )
        assert np.isnan(hourly.loc[gap_ts, "storage_pct"])

    def test_none_and_empty(self):
        assert hourly_frame(None) is None
        assert hourly_frame(pd.DataFrame()) is None


class TestLoadDaily:
    def test_empty_db_result_does_not_wipe_cache(self, tmp_path, pg_url):
        from sqlalchemy import create_engine, text

        from src.ml.dam import CACHE_FILE, load_daily

        # Good cache from a previous run
        cache_path = tmp_path / CACHE_FILE
        _daily(3).rename_axis("date").to_csv(cache_path, compression="gzip")
        # Reachable DB whose table exists but is empty
        engine = create_engine(pg_url)
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE rid_reservoir_daily (dam_id TEXT, date DATE, "
                    "storage_pct REAL, inflow_mcm REAL, outflow_mcm REAL)"
                )
            )
        engine.dispose()
        result = load_daily(db_url=pg_url, cache_dir=tmp_path)
        assert result.empty  # honest empty result...
        cached = pd.read_csv(cache_path, index_col=0)
        assert len(cached) == 3  # ...but the good cache survives


def _grid(hours=400, start="2024-09-15"):
    idx = pd.date_range(start, periods=hours, freq="h")
    frames = []
    for code in ("P.1", "P.82"):
        frames.append(
            pd.DataFrame(
                {
                    "timestamp": idx,
                    "station_code": code,
                    "water_level": 2.0,
                    "discharge": 100.0,
                }
            )
        )
    return features.make_hourly_grid(pd.concat(frames, ignore_index=True))


class TestFeatureGating:
    def test_dam_columns_only_for_dam_stations(self):
        grid = _grid()
        dam = hourly_frame(_daily(days=20, start="2024-09-10"))
        X_p1 = features.build_features(grid, "P.1", dam=dam)
        X_p82 = features.build_features(grid, "P.82", dam=dam)
        for col in features.DAM_FEATURES:
            assert col in X_p1.columns
            assert col not in X_p82.columns
        # values actually aligned, not all-NaN
        assert X_p1["dam_storage_pct"].notna().any()
        assert X_p1["dam_inflow"].notna().any()

    def test_none_dam_omits_columns(self):
        X = features.build_features(_grid(), "P.1", dam=None)
        for col in features.DAM_FEATURES:
            assert col not in X.columns

    def test_empty_dam_frame_yields_nan_columns(self):
        # Serving contract: empty frame -> columns exist as NaN so dam-trained
        # bundles pass the feature guard when the DB read fails.
        X = features.build_features(_grid(), "P.1", dam=pd.DataFrame())
        for col in features.DAM_FEATURES:
            assert col in X.columns
            assert X[col].isna().all()

    def test_storage_delta_72h(self):
        grid = _grid(hours=24 * 12, start="2024-09-15")
        dam = hourly_frame(_daily(days=20, start="2024-09-10"))
        X = features.build_features(grid, "P.1", dam=dam)
        ts = pd.Timestamp("2024-09-24 12:00")
        expected = (
            X.loc[ts, "dam_storage_pct"]
            - X.loc[ts - pd.Timedelta(hours=72), "dam_storage_pct"]
        )
        assert abs(X.loc[ts, "dam_storage_pct_d3"] - expected) < 1e-9