"""rain_fc24 at the newest grid row (the one serving predicts from)."""

import numpy as np
import pandas as pd

from src.ml import features
from src.ml.data import _prepare_hii_levels


def _grid(hours=48):
    idx = pd.date_range("2026-09-01", periods=hours, freq="h")
    df = pd.DataFrame(
        {
            "timestamp": idx,
            "station_code": "P.1",
            "water_level": np.linspace(1.0, 2.0, hours),
            "discharge": np.nan,
        }
    )
    return features.make_hourly_grid(df)


def test_forward_rain_reaches_past_the_last_reading():
    grid = _grid()
    idx = grid.observed.index
    rain = pd.Series(
        np.arange(1, len(idx) + 49, dtype=float),
        index=pd.date_range(idx[0], periods=len(idx) + 48, freq="h"),
    )
    X = features.build_features(grid, "P.1", rain=rain)
    last = idx[-1]
    expected = rain.loc[last + pd.Timedelta(hours=1) : last + pd.Timedelta(hours=24)]
    assert X.loc[last, "rain_fc24"] == expected.sum()
    # rows whose window lies inside the grid are what they always were
    old = rain.reindex(idx).shift(-1).iloc[::-1].rolling(24, min_periods=1).sum()
    inner = idx[:-24]
    np.testing.assert_allclose(X.loc[inner, "rain_fc24"], old.iloc[::-1].loc[inner])


def test_forward_rain_nan_without_forecast_hours():
    grid = _grid()
    idx = grid.observed.index
    rain = pd.Series(1.0, index=idx)  # nothing after the last reading
    X = features.build_features(grid, "P.1", rain=rain)
    assert np.isnan(X.loc[idx[-1], "rain_fc24"])
    assert X.loc[idx[-2], "rain_fc24"] == 1.0


def test_hii_rows_latest_reading_in_an_hour_wins():
    raw = pd.DataFrame(
        {
            "timestamp": pd.to_datetime(
                ["2026-10-01 10:40", "2026-10-01 10:10", "2026-10-01 11:00"]
            ),
            "station_code": "P.1",
            "wl_msl": [301.5, 301.2, 301.7],
            "discharge": [10.0, 9.0, 11.0],
        }
    )
    out = _prepare_hii_levels(raw)
    assert list(out["wl_msl"]) == [301.5, 301.7]
    assert list(out["timestamp"]) == list(
        pd.to_datetime(["2026-10-01 10:00", "2026-10-01 11:00"])
    )