openrun/tests/unit/test_banister.py

"""Tests for Banister CTL/ATL/TSB and the forecast splice helper.

Two things to lock down:
  - the EWMA math itself, via a closed-form impulse-response check.
  - the forecast splice, which must produce the same output as feeding the
    whole series at once (i.e. no double-counting on the boundary day).
"""

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest

from openrun.model import banister, banister_forecast


def _daily_zero_load(start: str, days: int) -> pd.Series:
    idx = pd.date_range(start, periods=days, freq="D")
    return pd.Series(0.0, index=idx)


# ---------------------------------------------------------------------------
# banister() — closed-form correctness
# ---------------------------------------------------------------------------

def test_banister_impulse_decays_exponentially() -> None:
    """Single load on day 0, zeros after. The recursion
        CTL[i] = CTL[i-1]*decay + load[i]*(1-decay)
    with decay = exp(-1/tau) yields
        CTL[i] = L * (1-decay) * decay**i.
    """
    load = _daily_zero_load("2026-01-01", 50)
    load.iloc[0] = 100.0

    pmc = banister(load, ctl_tau=42.0, atl_tau=7.0)
    decay_ctl = math.exp(-1 / 42)
    decay_atl = math.exp(-1 / 7)

    for i in (0, 1, 7, 21, 42):
        expected_ctl = 100.0 * (1 - decay_ctl) * decay_ctl**i
        expected_atl = 100.0 * (1 - decay_atl) * decay_atl**i
        assert pmc["CTL"].iloc[i] == pytest.approx(expected_ctl, rel=1e-9)
        assert pmc["ATL"].iloc[i] == pytest.approx(expected_atl, rel=1e-9)


def test_banister_steady_state_approaches_input() -> None:
    """Constant input L converges to CTL=L; after `tau` days CTL ≈ L·(1-1/e)."""
    load = pd.Series(100.0, index=pd.date_range("2026-01-01", periods=200, freq="D"))
    pmc = banister(load, ctl_tau=42.0, atl_tau=7.0)
    assert pmc["CTL"].iloc[41] == pytest.approx(100.0 * (1 - 1 / math.e), rel=1e-3)
    # After ~5τ days (200/42 ≈ 4.76), CTL is within 1% of the steady-state value.
    assert pmc["CTL"].iloc[-1] == pytest.approx(100.0, rel=0.01)


def test_banister_rest_days_filled_with_zero() -> None:
    """Sparse input (only some dates) must be reindexed daily with zeros so
    both EWMAs continue to decay on rest days."""
    idx = pd.to_datetime(["2026-01-01", "2026-01-10"])
    load = pd.Series([100.0, 0.0], index=idx)
    pmc = banister(load, ctl_tau=42.0, atl_tau=7.0)
    assert len(pmc) == 10  # 2026-01-01 through 2026-01-10 inclusive


def test_banister_empty_input_returns_empty_frame() -> None:
    out = banister(pd.Series(dtype=float))
    assert list(out.columns) == ["CTL", "ATL", "TSB"]
    assert len(out) == 0


# ---------------------------------------------------------------------------
# banister_forecast() — splice + invariant
# ---------------------------------------------------------------------------

def test_banister_forecast_matches_full_history() -> None:
    """Splitting a full series into (history, future) and running through the
    forecast helper produces the identical PMC frame to feeding the whole
    series at once. This is the load-bearing invariant."""
    rng = np.random.default_rng(seed=42)
    idx = pd.date_range("2026-01-01", periods=120, freq="D")
    full = pd.Series(rng.integers(0, 150, size=120).astype(float), index=idx)

    cutoff = idx[60]
    history = full[full.index <= cutoff]
    future = full[full.index > cutoff]

    pmc_full = banister(full)
    pmc_spliced = banister_forecast(history, future, today=cutoff)

    pd.testing.assert_frame_equal(pmc_full, pmc_spliced)


def test_banister_forecast_default_today_is_history_max() -> None:
    history = pd.Series([10.0, 20.0, 30.0],
                        index=pd.date_range("2026-01-01", periods=3, freq="D"))
    future = pd.Series([5.0, 5.0],
                       index=pd.date_range("2026-01-04", periods=2, freq="D"))

    out_default = banister_forecast(history, future)
    out_explicit = banister_forecast(history, future, today="2026-01-03")
    pd.testing.assert_frame_equal(out_default, out_explicit)


def test_banister_forecast_drops_future_dates_at_or_before_today() -> None:
    """A future row dated ≤ today must be dropped — history is authoritative for that day."""
    history = pd.Series([100.0],
                        index=pd.date_range("2026-01-03", periods=1, freq="D"))
    # Future contains an "overlapping" entry at today that would otherwise sum with history.
    future = pd.Series([999.0, 5.0],
                       index=pd.to_datetime(["2026-01-03", "2026-01-04"]))

    pmc = banister_forecast(history, future, today="2026-01-03")
    # Day 0 should reflect history's 100, not the 999 from future.
    decay_ctl = math.exp(-1 / 42)
    assert pmc["CTL"].iloc[0] == pytest.approx(100.0 * (1 - decay_ctl), rel=1e-9)


def test_banister_forecast_extends_through_future_end() -> None:
    history = pd.Series([10.0],
                        index=pd.date_range("2026-01-01", periods=1, freq="D"))
    future = pd.Series([5.0] * 5,
                       index=pd.date_range("2026-01-02", periods=5, freq="D"))
    pmc = banister_forecast(history, future)
    assert pmc.index.max() == pd.Timestamp("2026-01-06")