"""Tests for model_based.py.

The Wiener tests pin the Wiener-Hopf solve and check that the predicted minimum
MSE matches the MSE actually achieved on a long record. The Kalman tests pin the
filter against its two limits (trust-the-measurement, trust-the-model) and check
the headline self-consistency: the error covariance the filter reports really is
the covariance of its errors, confirmed by Monte Carlo.
"""

import numpy as np
from scipy.linalg import toeplitz

from model_based import (
    wiener_hopf,
    wiener_min_mse,
    wiener_fir_design,
    wiener_apply,
    wiener_freq,
    kalman_filter,
    cv_model,
    simulate_cv,
    steady_state_gain,
)


# --- Wiener ----------------------------------------------------------------

def test_wiener_hopf_solves_normal_equations():
    # Build a valid (positive-definite) symmetric Toeplitz R from an
    # autocorrelation sequence, pick an arbitrary rdx, and confirm the solver
    # returns the exact solution of R w = rdx.
    rxx = np.array([4.0, 1.5, 0.5, 0.1])
    rdx = np.array([3.0, 1.0, 0.2, 0.0])
    w = wiener_hopf(rxx, rdx)
    R = toeplitz(rxx)
    assert np.allclose(R @ w, rdx, atol=1e-12)


def test_wiener_min_mse_matches_measured_mse():
    # AR(1) desired signal d in white noise. With the TRUE correlations the
    # predicted J_min must match the MSE the filter actually achieves.
    rng = np.random.default_rng(0)
    a, sigma_e, sigma_v = 0.8, 1.0, 1.0
    M = 16

    # True correlations: var_d = sigma_e^2 / (1 - a^2); rdd[k] = var_d a^|k|.
    var_d = sigma_e ** 2 / (1 - a ** 2)
    lags = np.arange(M)
    rdd = var_d * a ** lags
    rvv = np.zeros(M)
    rvv[0] = sigma_v ** 2
    rxx = rdd + rvv                # d and v uncorrelated
    rdx = rdd                     # E[d[n] x[n-k]] = E[d[n] d[n-k]]

    w = wiener_hopf(rxx, rdx)
    j_min = wiener_min_mse(rxx, rdx, var_d)

    # Generate a long realisation and measure the achieved MSE.
    n = 200_000
    e = rng.standard_normal(n) * sigma_e
    d = np.zeros(n)
    for k in range(1, n):
        d[k] = a * d[k - 1] + e[k]
    x = d + rng.standard_normal(n) * sigma_v
    d_hat = wiener_apply(w, x)
    measured = np.mean((d[M:] - d_hat[M:]) ** 2)   # drop the startup transient

    assert np.isclose(measured, j_min, rtol=0.05)
    # And the filter genuinely helps: MMSE well below the input noise power.
    assert j_min < sigma_v ** 2


def test_wiener_freq_gain_is_between_zero_and_one():
    s_dd = np.array([10.0, 5.0, 0.1, 0.01])
    s_vv = np.array([0.1, 1.0, 5.0, 10.0])
    H = wiener_freq(s_dd, s_vv)
    assert np.all(H >= 0.0) and np.all(H <= 1.0)
    # Signal-dominated bin passes; noise-dominated bin is suppressed.
    assert H[0] > 0.95
    assert H[-1] < 0.05


def test_wiener_fir_design_beats_input():
    # Data-driven design (no true spectra) should still reduce the error.
    rng = np.random.default_rng(1)
    a, n, M = 0.85, 50_000, 12
    e = rng.standard_normal(n)
    d = np.zeros(n)
    for k in range(1, n):
        d[k] = a * d[k - 1] + e[k]
    x = d + rng.standard_normal(n)
    w, j_min = wiener_fir_design(x, d, M)
    d_hat = wiener_apply(w, x)
    out_mse = np.mean((d[M:] - d_hat[M:]) ** 2)
    in_mse = np.mean((d[M:] - x[M:]) ** 2)
    assert out_mse < in_mse
    assert np.isclose(out_mse, j_min, rtol=0.1)


def test_wiener_fir_design_recovers_pure_delay():
    # An asymmetric d/x relation pins the cross-correlation lag SIGN that the
    # symmetric denoising test above cannot see. If d[n] = x[n-2], the optimal
    # causal FIR filter is a pure delay by two taps, w = [0, 0, 1, 0, 0, 0].
    # A flipped rdx lag sign would instead try to predict x[n+2] and fail.
    rng = np.random.default_rng(5)
    x = rng.standard_normal(20_000)
    d = np.zeros_like(x)
    d[2:] = x[:-2]
    w, j_min = wiener_fir_design(x, d, 6)
    assert np.allclose(w, [0.0, 0.0, 1.0, 0.0, 0.0, 0.0], atol=0.02)
    assert j_min < 1e-2


# --- Kalman ----------------------------------------------------------------

def test_cv_model_matches_closed_form():
    dt, sa, sz = 0.1, 0.5, 2.0
    F, H, Q, R = cv_model(dt, sa, sz)
    assert np.allclose(F, [[1, dt], [0, 1]])
    assert np.allclose(H, [[1, 0]])
    assert np.isclose(Q[0, 0], sa ** 2 * dt ** 4 / 4)
    assert np.isclose(Q[0, 1], sa ** 2 * dt ** 3 / 2)
    assert np.isclose(Q[1, 1], sa ** 2 * dt ** 2)
    assert np.isclose(R[0, 0], sz ** 2)


def test_kalman_trusts_measurement_when_R_tiny():
    # With near-zero measurement noise the position estimate should track the
    # measurements almost exactly.
    rng = np.random.default_rng(2)
    F, H, Q, R = cv_model(1.0, 0.3, 1.0)
    R = np.array([[1e-8]])
    _, z = simulate_cv(F, H, Q, R, [0.0, 1.0], 50, rng)
    xs, _ = kalman_filter(z, F, H, Q, R, x0=np.array([0.0, 1.0]),
                          P0=np.eye(2))
    assert np.allclose(xs[5:, 0], z[5:, 0], atol=1e-3)


def test_kalman_ignores_measurement_when_R_huge():
    # With enormous measurement noise the update barely corrects the prediction;
    # the estimate is essentially open-loop dead reckoning from x0.
    F, H, Q, R = cv_model(1.0, 1e-6, 1e12)
    z = np.zeros((20, 1))   # measurements are nonsense but down-weighted to nothing
    x0 = np.array([0.0, 2.0])
    xs, _ = kalman_filter(z, F, H, Q, R, x0=x0, P0=1e-6 * np.eye(2))
    # Pure prediction would give position = 2 * k (constant velocity 2).
    predicted = 2.0 * np.arange(1, 21)
    assert np.allclose(xs[:, 0], predicted, atol=1e-2)


def test_truck_rmse_beats_measurement():
    rng = np.random.default_rng(3)
    F, H, Q, R = cv_model(1.0, 0.2, 5.0)
    xs_true, z = simulate_cv(F, H, Q, R, np.array([0.0, 1.0]), 200, rng)
    xs, _ = kalman_filter(z, F, H, Q, R, x0=np.array([0.0, 1.0]),
                          P0=np.eye(2))
    rmse_filtered = np.sqrt(np.mean((xs[:, 0] - xs_true[:, 0]) ** 2))
    rmse_meas = np.sqrt(np.mean((z[:, 0] - xs_true[:, 0]) ** 2))
    assert rmse_filtered < rmse_meas


def test_error_covariance_is_consistent():
    # The headline guarantee: the P the filter reports is the actual covariance
    # of its estimation error. Estimate that covariance by Monte Carlo and
    # compare to the steady-state P from a single filter run.
    F, H, Q, R = cv_model(1.0, 0.3, 3.0)
    steps = 40
    x0 = np.array([0.0, 1.0])
    P0 = np.eye(2)

    # One run to get the filter's reported steady-state covariance.
    rng = np.random.default_rng(10)
    _, z = simulate_cv(F, H, Q, R, x0, steps, rng)
    _, Ps = kalman_filter(z, F, H, Q, R, x0, P0)
    P_reported = Ps[-1]

    # Many runs to estimate the true error covariance at the final step.
    runs = 3000
    errors = np.empty((runs, 2))
    rng = np.random.default_rng(20)
    for r in range(runs):
        xs_true, z = simulate_cv(F, H, Q, R, x0, steps, rng)
        xs, _ = kalman_filter(z, F, H, Q, R, x0, P0)
        errors[r] = xs_true[-1] - xs[-1]
    P_empirical = np.cov(errors, rowvar=False)

    assert np.allclose(P_empirical, P_reported, rtol=0.15, atol=0.05)


def test_steady_state_gain_matches_running_filter():
    # The per-step covariance of a long filter run must converge to the
    # steady-state P computed from the Riccati recursion alone.
    rng = np.random.default_rng(4)
    F, H, Q, R = cv_model(1.0, 0.4, 2.0)
    K_inf, P_inf = steady_state_gain(F, H, Q, R)
    _, z = simulate_cv(F, H, Q, R, np.array([0.0, 1.0]), 300, rng)
    _, Ps = kalman_filter(z, F, H, Q, R, x0=np.array([0.0, 1.0]), P0=np.eye(2))
    assert np.allclose(Ps[-1], P_inf, atol=1e-6)