# =============================================================================
# QWHPI — Quebec Weekly Housing Price Index
# Author  : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# File    : engine/tests/test_state_space.py
# Purpose : Unit tests for the heteroskedastic local-level Kalman engine.
# =============================================================================
"""State-space unit tests."""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import numpy as np

from qwhpi.state_space import fit_local_level

RNG = np.random.default_rng(0)


def test_recovers_constant_signal():
    T, n = 150, 30
    y = 0.05 + RNG.normal(0, 0.2 / np.sqrt(n), T)
    fit = fit_local_level(y, np.full(T, n), sigma2=0.04)
    assert abs(fit.smoothed[T // 2] - 0.05) < 0.02
    assert fit.tau2 < 1e-4  # constant state -> tiny innovation variance


def test_tracks_drifting_signal():
    T, n = 200, 50
    truth = np.cumsum(RNG.normal(0, 0.01, T))
    y = truth + RNG.normal(0, 0.3 / np.sqrt(n), T)
    fit = fit_local_level(y, np.full(T, n), sigma2=0.09)
    rmse = np.sqrt(np.mean((fit.smoothed - truth) ** 2))
    assert rmse < 0.02


def test_empty_weeks_are_predicted_not_dropped():
    T = 100
    n = np.full(T, 20.0)
    n[40:60] = 0  # 20-week gap
    truth = np.cumsum(RNG.normal(0, 0.02, T))  # real innovation -> tau2 > 0
    y = np.where(n > 0, truth + RNG.normal(0, 0.05, T), np.nan)
    fit = fit_local_level(y, n, sigma2=0.05)
    # gap weeks still get an estimate; uncertainty grows through the gap
    assert np.all(np.isfinite(fit.smoothed))
    assert fit.filtered_var[59] > fit.filtered_var[41] > fit.filtered_var[39]
    assert fit.gain[50] == 0.0  # no local data -> fully parent/past-driven


def test_thin_weeks_shrink_more_than_liquid_weeks():
    T = 120
    n = np.full(T, 2.0)
    n[:60] = 200.0
    y = RNG.normal(0, 0.05, T)
    fit = fit_local_level(y, n, sigma2=0.09)
    assert fit.gain[:60].mean() > fit.gain[60:].mean()