gendesign/backend/tests/services/forecasting/test_regression.py

"""Unit tests for §9.6 Almon distributed-lag regression (Forgejo #978 Part B).

Covers the PURE numpy logic on SYNTHETIC series with a KNOWN injected lag effect:
  - _build_lag_matrix     — full-row-only lag profile, drops incomplete/None rows
  - _almon_basis          — W[j,p] = j^p (constrains 7 lags to degree+1 params)
  - newey_west_bandwidth  — floor(4·(n/100)^(2/9)) rule, ≥1 floor
  - newey_west_cov        — HAC covariance differs from naive OLS; PSD; manual NW
  - fit_almon_dl          — recovers the injected best_lag + sign + long-run; R²;
                            per-lag reconstruction; HAC SEs computed
  - build_fit_result      — gate (n≥30 ∧ R²≥0.1 ∧ Σβ<0) → regression vs fallback;
                            fallback on thin n / weak R² / wrong sign (no crash)
  - _build_phrase         — §9.6 text from the lag shape; insufficient on no-gate
  - compute_district_rate_regression — DB orchestrator wiring (mocked session)

NO live DB: the orchestrator test injects a fake session + monkeypatched data
loaders. Set a dummy DATABASE_URL BEFORE importing so app.core.config.Settings
fail-fast doesn't trip (same pattern as test_rate_sensitivity.py).
"""

from __future__ import annotations

import datetime as dt
import math
import os

os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost:5432/test")

import numpy as np
import pytest

from app.services.forecasting import regression as reg

# --------------------------------------------------------------------------- #
# Synthetic-series helpers — inject a KNOWN distributed-lag effect
# --------------------------------------------------------------------------- #


def _aperiodic_rate_deltas(n: int, *, seed: int = 13) -> list[float]:
    """Δrate series with APERIODIC (LCG) jitter → low autocorrelation across lags.

    A periodic regressor would let false lags compete with the injected one; an
    LCG jitter keeps successive Δ weakly correlated so the true lag shape wins.
    Finite from index 0 (the DL matrix builder drops incomplete leading rows).
    """
    lvl = 10.0
    state = seed
    levels: list[float] = []
    for _ in range(n):
        state = (state * 1103515245 + 12345) % 2147483648
        lvl += 0.3 + (state / 2147483648.0 - 0.5) * 0.8
        levels.append(lvl)
    return [0.0] + [levels[i] - levels[i - 1] for i in range(1, n)]


def _hump_beta(max_lag: int, *, peak: int, scale: float = 0.06) -> np.ndarray:
    """A negative 'hump' lag shape peaking (in magnitude) at ``peak``.

    |β_j| = scale − 0.012·(j−peak)² (floored at 0.005), all signs negative — the
    economically expected shape (rate ↑ → demand ↓, response builds then fades).
    Representable approximately by an Almon deg-2 polynomial, so the fit recovers
    the peak and long-run.
    """
    betas: list[float] = []
    for j in range(max_lag + 1):
        mag = scale - 0.012 * (j - peak) ** 2
        betas.append(-max(0.005, mag))
    return np.asarray(betas, dtype=float)


def _y_from_lag_shape(
    x: list[float], beta: np.ndarray, *, max_lag: int, noise: float = 0.0, seed: int = 0
) -> list[float | None]:
    """y[t] = Σ_j β_j·x[t−j] (+ optional gaussian noise); y[t<max_lag] = None.

    Builds the regressand carrying the injected distributed-lag effect exactly
    (plus noise). Leading months without a full lag profile → None (the builder
    drops them anyway).
    """
    rng = np.random.default_rng(seed)
    y: list[float | None] = [None] * len(x)
    for t in range(max_lag, len(x)):
        val = float(sum(beta[j] * x[t - j] for j in range(max_lag + 1)))
        if noise > 0.0:
            val += float(rng.normal(0.0, noise))
        y[t] = val
    return y


# --------------------------------------------------------------------------- #
# _build_lag_matrix
# --------------------------------------------------------------------------- #


class TestBuildLagMatrix:
    def test_shapes_and_full_rows_only(self) -> None:
        x = [float(i) for i in range(10)]
        y = [float(i) * 0.1 for i in range(10)]
        built = reg._build_lag_matrix(x, y, max_lag=2)
        assert built is not None
        xm, yv = built
        # First usable row is t=max_lag=2 → 10−2 = 8 rows, 3 lag columns.
        assert xm.shape == (8, 3)
        assert yv.shape == (8,)
        # Row 0 corresponds to t=2: [x[2], x[1], x[0]] = [2,1,0].
        assert list(xm[0]) == [2.0, 1.0, 0.0]

    def test_drops_rows_with_none_in_any_lag(self) -> None:
        x: list[float | None] = [0.0, 1.0, None, 3.0, 4.0, 5.0]
        y: list[float | None] = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5]
        built = reg._build_lag_matrix(x, y, max_lag=2)
        assert built is not None
        xm, _yv = built
        # t=2 reads x[0..2] (has None) → dropped; t=3 reads x[1..3] (has None) →
        # dropped; t=4 reads x[2..4] (has None) → dropped; t=5 reads x[3..5] OK.
        assert xm.shape == (1, 3)
        assert list(xm[0]) == [5.0, 4.0, 3.0]

    def test_drops_rows_with_none_y(self) -> None:
        x = [float(i) for i in range(6)]
        y: list[float | None] = [0.0, 0.1, None, 0.3, 0.4, 0.5]
        built = reg._build_lag_matrix(x, y, max_lag=1)
        assert built is not None
        xm, yv = built
        # t=2 has y=None → dropped. Usable t ∈ {1,3,4,5} → 4 rows.
        assert xm.shape == (4, 2)
        assert yv.shape == (4,)

    def test_returns_none_when_no_full_row(self) -> None:
        x: list[float | None] = [None, None, None]
        y: list[float | None] = [1.0, 2.0, 3.0]
        assert reg._build_lag_matrix(x, y, max_lag=1) is None


# --------------------------------------------------------------------------- #
# _almon_basis
# --------------------------------------------------------------------------- #


class TestAlmonBasis:
    def test_j_to_the_p(self) -> None:
        w = reg._almon_basis(3, 2)  # lags 0..3, degree 2
        assert w.shape == (4, 3)
        # Column p = j^p: col0 = ones, col1 = j, col2 = j².
        assert list(w[:, 0]) == [1.0, 1.0, 1.0, 1.0]
        assert list(w[:, 1]) == [0.0, 1.0, 2.0, 3.0]
        assert list(w[:, 2]) == [0.0, 1.0, 4.0, 9.0]

    def test_reconstruct_quadratic_beta_exactly(self) -> None:
        # β_j = 2 − 0.5j + 0.1j² is degree-2 → W @ γ reproduces it for γ=[2,−0.5,0.1].
        w = reg._almon_basis(6, 2)
        gamma = np.array([2.0, -0.5, 0.1])
        beta = w @ gamma
        expected = np.array([2.0 - 0.5 * j + 0.1 * j**2 for j in range(7)])
        assert np.allclose(beta, expected)


# --------------------------------------------------------------------------- #
# newey_west_bandwidth / newey_west_cov
# --------------------------------------------------------------------------- #


class TestNeweyWestBandwidth:
    def test_rule_values(self) -> None:
        # floor(4·(n/100)^(2/9)).
        assert reg.newey_west_bandwidth(100) == 4
        assert reg.newey_west_bandwidth(41) == 3
        assert reg.newey_west_bandwidth(50) == 3

    def test_small_n_values(self) -> None:
        # floor(4·(10/100)^(2/9)) = floor(2.398) = 2; n=20 → floor(2.40)·… = 2.
        assert reg.newey_west_bandwidth(10) == 2
        assert reg.newey_west_bandwidth(20) == 2

    def test_floor_at_one_for_tiny_n(self) -> None:
        # n=2 → floor(4·0.02^(2/9)) ≈ floor(1.06) = 1, but the ≥1 floor guarantees
        # at least a lag-1 autocovariance whenever n>1.
        assert reg.newey_west_bandwidth(2) == 1
        assert reg.newey_west_bandwidth(3) == 1

    def test_zero_for_degenerate(self) -> None:
        assert reg.newey_west_bandwidth(1) == 0
        assert reg.newey_west_bandwidth(0) == 0


class TestNeweyWestCov:
    def test_psd_and_symmetric(self) -> None:
        rng = np.random.default_rng(1)
        n = 50
        design = np.column_stack([np.ones(n), rng.normal(size=(n, 2))])
        resid = rng.normal(size=n)
        cov = reg.newey_west_cov(design, resid, bandwidth=4)
        # Symmetric and positive semi-definite (Bartlett weights guarantee PSD).
        assert np.allclose(cov, cov.T, atol=1e-10)
        eig = np.linalg.eigvalsh(cov)
        assert float(eig.min()) >= -1e-8

    def test_bandwidth_zero_equals_white_hc0(self) -> None:
        rng = np.random.default_rng(2)
        n = 40
        design = np.column_stack([np.ones(n), rng.normal(size=(n, 1))])
        resid = rng.normal(size=n)
        cov0 = reg.newey_west_cov(design, resid, bandwidth=0)
        # HC0: (X'X)^-1 (Σ u² x x') (X'X)^-1 — reconstruct manually.
        xtx_inv = np.linalg.inv(design.T @ design)
        ux = design * resid.reshape(-1, 1)
        hc0 = xtx_inv @ (ux.T @ ux) @ xtx_inv
        assert np.allclose(cov0, hc0, atol=1e-12)

    def test_hac_differs_from_naive_under_autocorrelation(self) -> None:
        # Construct strongly AUTOCORRELATED residuals → HAC SE must differ from
        # the naive iid OLS SE (the whole point of NW).
        n = 80
        rng = np.random.default_rng(3)
        x = rng.normal(size=n)
        design = np.column_stack([np.ones(n), x])
        # AR(1) residuals (ρ=0.7) → positive autocorrelation.
        e = np.zeros(n)
        for t in range(1, n):
            e[t] = 0.7 * e[t - 1] + rng.normal(0, 1)
        hac = reg.newey_west_cov(design, e, bandwidth=reg.newey_west_bandwidth(n))
        sigma2 = float(e @ e) / (n - design.shape[1])
        naive = sigma2 * np.linalg.inv(design.T @ design)
        # The slope variance estimates must differ materially under autocorrelation.
        assert not math.isclose(hac[1, 1], naive[1, 1], rel_tol=0.05)


# --------------------------------------------------------------------------- #
# fit_almon_dl — recover the injected lag shape
# --------------------------------------------------------------------------- #


class TestFitAlmonDl:
    def test_recovers_injected_best_lag_and_sign(self) -> None:
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n, seed=13)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.002, seed=0)
        fit = reg.fit_almon_dl(x, y, max_lag=max_lag, degree=2)
        assert fit is not None
        # Peak lag recovered.
        assert fit["best_lag"] == 2
        # Long-run sign negative (rate ↑ → demand ↓) and close to the truth.
        assert fit["long_run_coef"] < 0
        assert math.isclose(fit["long_run_coef"], float(beta.sum()), abs_tol=0.02)
        # Clean injected signal → high R².
        assert fit["r2"] is not None and fit["r2"] > 0.8
        # Gate flags all green on this clean, long, correctly-signed series.
        assert fit["gate_n_ok"] and fit["gate_r2_ok"] and fit["gate_sign_ok"]

    def test_recovers_different_peak_lag(self) -> None:
        # Shift the injected peak to lag 4 → the fit must track it.
        n, max_lag = 64, 6
        x = _aperiodic_rate_deltas(n, seed=21)
        beta = _hump_beta(max_lag, peak=4)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.002, seed=1)
        fit = reg.fit_almon_dl(x, y, max_lag=max_lag, degree=2)
        assert fit is not None
        assert fit["best_lag"] == 4
        assert fit["long_run_coef"] < 0

    def test_per_lag_reconstruction_length_and_finite(self) -> None:
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.001, seed=2)
        fit = reg.fit_almon_dl(x, y, max_lag=max_lag, degree=2)
        assert fit is not None
        per_lag = fit["per_lag_coef"]
        assert len(per_lag) == max_lag + 1
        assert all(math.isfinite(c) for c in per_lag)

    def test_hac_se_computed_for_every_lag(self) -> None:
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.003, seed=3)
        fit = reg.fit_almon_dl(x, y, max_lag=max_lag, degree=2)
        assert fit is not None
        hac_se = fit["hac_se"]
        # One HAC SE per reconstructed lag coefficient, all finite and ≥0.
        assert len(hac_se) == max_lag + 1
        assert all(math.isfinite(s) and s >= 0.0 for s in hac_se)
        # Bandwidth follows the NW rule for this n.
        assert fit["hac_bandwidth"] == reg.newey_west_bandwidth(fit["n"])

    def test_degree_must_be_below_lag_count(self) -> None:
        # degree ≥ max_lag+1 is not a constraint (degenerates to free lags) → refuse.
        x = _aperiodic_rate_deltas(40)
        beta = _hump_beta(6, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=6, noise=0.0)
        assert reg.fit_almon_dl(x, y, max_lag=6, degree=7) is None

    def test_thin_series_returns_none(self) -> None:
        # Too few full rows to fit (< _MIN_FIT_OBS) → None, not a crash.
        x = _aperiodic_rate_deltas(10)
        beta = _hump_beta(6, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=6, noise=0.0)
        assert reg.fit_almon_dl(x, y, max_lag=6, degree=2) is None

    def test_zero_variance_y_returns_none(self) -> None:
        x = _aperiodic_rate_deltas(50)
        y: list[float | None] = [None] * 6 + [0.0] * 44  # constant → no variance
        assert reg.fit_almon_dl(x, y, max_lag=6, degree=2) is None


# --------------------------------------------------------------------------- #
# build_fit_result — gate (mirror _elasticity_coef) → regression vs fallback
# --------------------------------------------------------------------------- #


_SEG: dict[str, str | None] = {"district": "Академический", "obj_class": None}


class TestBuildFitResult:
    def test_gate_pass_emits_regression(self) -> None:
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n, seed=13)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.002, seed=0)
        res = reg.build_fit_result(x, y, segment=_SEG, max_lag=max_lag, degree=2)
        assert res.source == "regression"
        assert res.best_lag_months == 2
        assert res.coef is not None and res.coef < 0
        assert res.x_pct is not None and res.x_pct < 0  # demand drops
        assert res.r2 is not None and res.r2 > 0.8
        assert res.per_lag_coef is not None and len(res.per_lag_coef) == max_lag + 1
        assert res.hac_se is not None and len(res.hac_se) == max_lag + 1
        # Phrase carries the magnitude + peak lag.
        assert "снижается" in res.phrase
        assert f"{abs(round(res.x_pct, 1))}" in res.phrase

    def test_thin_n_degrades_to_fallback(self) -> None:
        # Enough to fit, but n < _MIN_OBS (30) → gate fails on n → fallback. We keep
        # the diagnostic numbers (per_lag/r2/n) but make no claim.
        n, max_lag = 28, 6  # ~22 usable rows < 30
        x = _aperiodic_rate_deltas(n, seed=5)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.001, seed=4)
        res = reg.build_fit_result(x, y, segment=_SEG, max_lag=max_lag, degree=2)
        assert res.source == "fallback"
        assert res.n < reg._MIN_OBS
        assert res.coef is None and res.x_pct is None and res.best_lag_months is None
        assert res.phrase == reg._PHRASE_INSUFFICIENT
        # Diagnostics retained (mirror _elasticity_coef returning r2/n in fallback).
        assert res.per_lag_coef is not None

    def test_wrong_sign_degrades_to_fallback(self) -> None:
        # POSITIVE long-run (rate ↑ → demand ↑) violates the gate sign → fallback,
        # even with plenty of obs and a strong fit.
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n, seed=13)
        beta = -_hump_beta(max_lag, peak=2)  # flip all signs → positive long-run
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.002, seed=0)
        res = reg.build_fit_result(x, y, segment=_SEG, max_lag=max_lag, degree=2)
        assert res.source == "fallback"
        assert res.coef is None
        assert res.phrase == reg._PHRASE_INSUFFICIENT

    def test_weak_r2_degrades_to_fallback(self) -> None:
        # Pure noise regressand (no rate link) at large n: a 3-param Almon basis
        # cannot overfit ~114 noise points, so R² collapses well below 0.1 → the
        # gate fails on R² (or sign) → fallback. (At small n a flexible basis can
        # spuriously clear R²≥0.1 — which is exactly why the n≥30 gate + advisory
        # status exist; here we use n=120 so the no-signal case is unambiguous.)
        n, max_lag = 120, 6
        x = _aperiodic_rate_deltas(n, seed=13)
        rng = np.random.default_rng(7)
        y: list[float | None] = [None] * max_lag + [
            float(v) for v in rng.normal(0, 0.05, size=n - max_lag)
        ]
        res = reg.build_fit_result(x, y, segment=_SEG, max_lag=max_lag, degree=2)
        assert res.source == "fallback"
        assert res.coef is None
        # Confirm it degraded specifically because the fit explains ~no variance.
        assert res.r2 is not None and res.r2 < reg._MIN_R2

    def test_empty_series_is_fallback_not_crash(self) -> None:
        res = reg.build_fit_result([], [], segment=_SEG)
        assert res.source == "fallback"
        assert res.n == 0
        assert res.phrase == reg._PHRASE_INSUFFICIENT

    def test_as_dict_shape(self) -> None:
        n, max_lag = 60, 6
        x = _aperiodic_rate_deltas(n, seed=13)
        beta = _hump_beta(max_lag, peak=2)
        y = _y_from_lag_shape(x, beta, max_lag=max_lag, noise=0.002, seed=0)
        d = reg.build_fit_result(x, y, segment=_SEG, max_lag=max_lag, degree=2).as_dict()
        for key in (
            "segment",
            "best_lag_months",
            "coef",
            "x_pct",
            "r2",
            "n",
            "per_lag_coef",
            "hac_se",
            "hac_bandwidth",
            "almon_degree",
            "source",
            "phrase",
        ):
            assert key in d
        assert d["source"] == "regression"
        assert isinstance(d["per_lag_coef"], list)


# --------------------------------------------------------------------------- #
# _build_phrase
# --------------------------------------------------------------------------- #


class TestBuildPhrase:
    def test_phrase_from_shape(self) -> None:
        p = reg._build_phrase(x_pct=-3.2, best_lag=2, gated=True)
        assert "3.2%" in p
        assert "2 мес" in p
        assert "снижается" in p

    def test_insufficient_when_not_gated(self) -> None:
        assert reg._build_phrase(x_pct=-3.2, best_lag=2, gated=False) == reg._PHRASE_INSUFFICIENT

    def test_insufficient_when_none(self) -> None:
        assert reg._build_phrase(x_pct=None, best_lag=2, gated=True) == reg._PHRASE_INSUFFICIENT
        assert reg._build_phrase(x_pct=-3.2, best_lag=None, gated=True) == reg._PHRASE_INSUFFICIENT


# --------------------------------------------------------------------------- #
# compute_district_rate_regression — DB orchestrator (mocked)
# --------------------------------------------------------------------------- #


class _FakeMacro:
    def __init__(self, month: dt.date, key_rate: float | None) -> None:
        self.month = month
        self.key_rate = key_rate


class _FakeSales:
    def __init__(self, months: list[dt.date], units: list[int]) -> None:
        self.months = months
        self.units = units


def _months(n: int) -> list[dt.date]:
    out: list[dt.date] = []
    y, m = 2021, 1
    for _ in range(n):
        out.append(dt.date(y, m, 1))
        m += 1
        if m == 13:
            m = 1
            y += 1
    return out


class TestComputeDistrictRateRegression:
    def test_orchestrator_wires_macro_and_sales(self, monkeypatch: pytest.MonkeyPatch) -> None:
        # Build a macro key_rate series whose Δ drives a lag-2 demand response, then
        # confirm the orchestrator assembles X (Δrate) and Y (Δln units), aligns
        # them, and recovers the injected lag via the pure fit. (The orchestrator
        # uses the module-default max_lag=6 internally.)
        n = 60
        months = _months(n)
        # key_rate levels: integrate the aperiodic Δ so _delta() recovers them.
        xdelta = _aperiodic_rate_deltas(n, seed=13)
        levels: list[float] = []
        acc = 10.0
        for d in xdelta:
            acc += d
            levels.append(acc)
        macro = [_FakeMacro(months[i], levels[i]) for i in range(n)]

        # Units carrying the lag-2 signal: ln(u_t) = ln(base) + Σ_{k≤t} β·Δrate[k-lag].
        beta_scalar = -0.05
        lag = 2
        ln_u = math.log(1000.0)
        units: list[int] = []
        for t in range(n):
            if t > 0:
                src = xdelta[t - lag] if t - lag >= 0 else 0.0
                ln_u += beta_scalar * src
            units.append(max(1, round(math.exp(ln_u))))
        sales = _FakeSales(months, units)

        monkeypatch.setattr(reg, "get_monthly_macro", lambda db, months_back: macro)
        monkeypatch.setattr(reg, "build_sales_series", lambda db, spec, source, months_back: sales)

        res = reg.compute_district_rate_regression(
            object(),  # type: ignore[arg-type]
            district="Академический",
            months_back=n,
        )
        assert res.segment["district"] == "Академический"
        assert res.source == "regression"
        # The single-lag injection at lag 2 → Almon shape peaks near lag 2.
        assert res.best_lag_months in (1, 2, 3)
        assert res.coef is not None and res.coef < 0
        assert res.n >= reg._MIN_OBS

    def test_orchestrator_graceful_on_empty(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(reg, "get_monthly_macro", lambda db, months_back: [])
        monkeypatch.setattr(
            reg,
            "build_sales_series",
            lambda db, spec, source, months_back: _FakeSales([], []),
        )
        res = reg.compute_district_rate_regression(
            object(),  # type: ignore[arg-type]
            district="Пустой",
        )
        assert res.source == "fallback"
        assert res.phrase == reg._PHRASE_INSUFFICIENT