"""Tests for recursive.py.

The two identities the topic page leans on are pinned at machine precision:
RLS at lambda = 1 equals batch (ridge) least squares, and the Kalman filter
on a static parameter equals recursive WLS.  The standing arc rules are
honored: the estimator covariance is checked against a numerically built
Fisher information matrix, and every claimed relationship between quantities
is asserted as a relationship, not just via its operands.
"""

import numpy as np
import pytest

from recursive import (attenuation_factor, batch_least_squares, block_rls,
                       constrained_least_squares, ew_variance_factor,
                       fd_tls_criterion, kalman_parameter_filter,
                       RecursiveLeastSquares, total_least_squares,
                       variance_equivalent_window)


def make_problem(rng, n=200, p=4, sigma=0.1, hetero=False):
    X = rng.normal(size=(n, p))
    theta = rng.normal(size=p)
    if hetero:
        noise_var = rng.uniform(0.04, 4.0, size=n)
    else:
        noise_var = np.full(n, sigma**2)
    y = X @ theta + np.sqrt(noise_var) * rng.normal(size=n)
    return X, y, theta, noise_var


# --------------------------------------------------------------- batch LS

def test_batch_ls_recovers_parameters():
    rng = np.random.default_rng(1)
    X, y, theta, _ = make_problem(rng, n=2000, sigma=0.05)
    est, _ = batch_least_squares(X, y)
    assert np.allclose(est, theta, atol=0.02)


def test_batch_wls_matches_normal_equations():
    rng = np.random.default_rng(2)
    X, y, _, noise_var = make_problem(rng, hetero=True)
    est, P = batch_least_squares(X, y, noise_var=noise_var)
    W = np.diag(1.0 / noise_var)
    direct = np.linalg.solve(X.T @ W @ X, X.T @ W @ y)
    assert np.allclose(est, direct, atol=1e-10)
    assert np.allclose(P, np.linalg.inv(X.T @ W @ X), atol=1e-12)


def test_batch_ls_covariance_matches_numeric_fisher_information():
    """Standing rule: any claimed bound is checked against a numeric FIM.

    For the linear Gaussian model the covariance sigma^2 (X^T X)^-1 IS the
    CRLB; build the Fisher information numerically from the log-likelihood
    curvature and compare.
    """
    rng = np.random.default_rng(3)
    sigma = 0.3
    X, y, theta, _ = make_problem(rng, n=50, p=3, sigma=sigma)

    def loglik(t):
        r = X @ theta - X @ t          # noiseless residual: E over data
        return -0.5 * np.sum(r**2) / sigma**2

    h = 1e-5
    p = len(theta)
    F_num = np.zeros((p, p))
    for i in range(p):
        for j in range(p):
            ei = np.zeros(p); ei[i] = h
            ej = np.zeros(p); ej[j] = h
            F_num[i, j] = -(loglik(theta + ei + ej) - loglik(theta + ei - ej)
                            - loglik(theta - ei + ej)
                            + loglik(theta - ei - ej)) / (4 * h * h)
    F_closed = X.T @ X / sigma**2
    assert np.allclose(F_num, F_closed, rtol=1e-4)
    _, P = batch_least_squares(X, y)
    assert np.allclose(sigma**2 * P, np.linalg.inv(F_num), rtol=1e-3)


def test_batch_ls_input_validation():
    with pytest.raises(ValueError):
        batch_least_squares(np.ones((3, 5)), np.ones(3))          # underdetermined
    with pytest.raises(ValueError):
        batch_least_squares(np.ones((5, 2)), np.ones(4))          # shape mismatch
    with pytest.raises(ValueError):
        batch_least_squares(np.ones((5, 2)), np.ones(5), noise_var=-1.0)
    with pytest.raises(ValueError):
        batch_least_squares(np.ones((5, 2)), np.ones(5), ridge=-0.1)


# --------------------------------------------------- RLS batch equivalence

def test_rls_lambda_one_equals_batch_ls():
    """THE identity of the page: lambda = 1, vanishing delta -> batch LS."""
    rng = np.random.default_rng(4)
    X, y, _, _ = make_problem(rng)
    rls = RecursiveLeastSquares(X.shape[1], forgetting=1.0, delta=1e-10)
    rls.fit(X, y)
    batch, _ = batch_least_squares(X, y)
    assert np.allclose(rls.theta, batch, atol=1e-7)


def test_rls_finite_delta_equals_ridge_exactly():
    """With finite delta the equivalence is EXACT: RLS(delta) == ridge(delta)."""
    rng = np.random.default_rng(5)
    X, y, _, _ = make_problem(rng)
    for delta in (1.0, 0.1, 17.3):
        rls = RecursiveLeastSquares(X.shape[1], forgetting=1.0, delta=delta)
        rls.fit(X, y)
        ridge, P_ridge = batch_least_squares(X, y, ridge=delta)
        assert np.allclose(rls.theta, ridge, atol=1e-12)
        assert np.allclose(rls.P, P_ridge, atol=1e-12)


def test_rls_P_is_inverse_normal_matrix():
    rng = np.random.default_rng(6)
    X, y, _, _ = make_problem(rng, n=60)
    delta = 2.5
    rls = RecursiveLeastSquares(X.shape[1], forgetting=1.0, delta=delta)
    rls.fit(X, y)
    direct = np.linalg.inv(delta * np.eye(X.shape[1]) + X.T @ X)
    assert np.allclose(rls.P, direct, atol=1e-12)


def test_rls_update_returns_a_priori_error():
    rls = RecursiveLeastSquares(2, delta=1e-6)
    e = rls.update(np.array([1.0, 0.0]), 3.0)
    assert e == pytest.approx(3.0)                     # theta started at zero
    e2 = rls.update(np.array([1.0, 0.0]), 3.0)
    assert abs(e2) < abs(e)                            # estimate moved toward y


def test_rls_input_validation():
    with pytest.raises(ValueError):
        RecursiveLeastSquares(2, forgetting=0.0)
    with pytest.raises(ValueError):
        RecursiveLeastSquares(2, forgetting=1.2)
    with pytest.raises(ValueError):
        RecursiveLeastSquares(2, delta=0.0)
    with pytest.raises(ValueError):
        RecursiveLeastSquares(2, theta0=np.zeros(3))


# --------------------------------------------------------------- block RLS

def test_block_rls_equals_batch():
    rng = np.random.default_rng(7)
    X, y, _, _ = make_problem(rng)
    for B in (8, 25, 50):
        theta, _ = block_rls(X, y, block_size=B, delta=1e-8)
        batch, _ = batch_least_squares(X, y)
        assert np.allclose(theta, batch, atol=1e-6), f"B={B}"


def test_block_rls_block_one_equals_sample_rls():
    rng = np.random.default_rng(8)
    X, y, _, _ = make_problem(rng, n=80)
    theta_b, P_b = block_rls(X, y, block_size=1, delta=1e-6)
    rls = RecursiveLeastSquares(X.shape[1], delta=1e-6)
    rls.fit(X, y)
    assert np.allclose(theta_b, rls.theta, atol=1e-9)
    assert np.allclose(P_b, rls.P, atol=1e-9)


def test_block_rls_finite_delta_equals_ridge():
    """The solve-based block update hits the ridge target at machine precision."""
    rng = np.random.default_rng(9)
    X, y, _, _ = make_problem(rng)
    delta = 1e-4
    theta, _ = block_rls(X, y, block_size=25, delta=delta)
    ridge, _ = batch_least_squares(X, y, ridge=delta)
    assert np.allclose(theta, ridge, atol=1e-11)


def test_block_rls_validation():
    X = np.ones((10, 2)); y = np.ones(10)
    with pytest.raises(ValueError):
        block_rls(X, y, block_size=0)
    with pytest.raises(ValueError):
        block_rls(X, y, block_size=11)
    with pytest.raises(ValueError):
        block_rls(X, y, block_size=2, delta=0.0)


# ------------------------------------------------------------ forgetting

def test_ew_variance_factor_monte_carlo():
    """MC variance of the EW mean matches sigma^2 (1-lam)/(1+lam)."""
    rng = np.random.default_rng(10)
    lam = 0.9
    n_mc, n_len = 1500, 600
    est = np.empty(n_mc)
    for m in range(n_mc):
        rls = RecursiveLeastSquares(1, forgetting=lam, delta=1e-8)
        rls.fit(np.ones((n_len, 1)), rng.normal(size=n_len))
        est[m] = rls.theta[0]
    assert np.var(est) == pytest.approx(ew_variance_factor(lam), rel=0.12)


def test_window_conventions_differ_by_factor_two():
    """The relationship itself: variance window ~ 2x the weight-sum window."""
    for lam in (0.9, 0.98, 0.995):
        weight_sum = 1.0 / (1.0 - lam)
        var_eq = variance_equivalent_window(lam)
        assert var_eq == pytest.approx((1 + lam) / (1 - lam), rel=1e-12)
        # ratio -> 2 as lam -> 1; already within 6% at lam = 0.9
        assert var_eq / weight_sum == pytest.approx(1.0 + lam, rel=1e-12)
        assert abs(var_eq / weight_sum - 2.0) < 0.11


def test_variance_equivalent_window_matches_plain_mean():
    """A plain mean over N_eff samples has the same variance as the EW mean."""
    lam = 0.95
    n_eff = variance_equivalent_window(lam)
    assert ew_variance_factor(lam) == pytest.approx(1.0 / n_eff, rel=1e-12)


def test_forgetting_validation():
    with pytest.raises(ValueError):
        ew_variance_factor(1.0)
    with pytest.raises(ValueError):
        variance_equivalent_window(0.0)


# ------------------------------------------------------- Kalman connection

def test_kalman_static_equals_rls():
    """Kalman with F=I, Q=0, R=1, P0=I/delta reproduces RLS trajectories."""
    rng = np.random.default_rng(11)
    X, y, _, _ = make_problem(rng)
    delta = 1e-6
    traj_k, P_k = kalman_parameter_filter(X, y, meas_var=1.0,
                                          P0=1.0 / delta)
    rls = RecursiveLeastSquares(X.shape[1], delta=delta)
    traj_r = rls.fit(X, y)
    assert np.allclose(traj_k, traj_r, atol=1e-9)
    assert np.allclose(P_k, rls.P, atol=1e-9)


def test_kalman_static_equals_batch_wls():
    """THE second identity: static-parameter Kalman == recursive WLS == batch WLS."""
    rng = np.random.default_rng(12)
    X, y, _, noise_var = make_problem(rng, hetero=True)
    traj, _ = kalman_parameter_filter(X, y, meas_var=noise_var, P0=1e10)
    wls, _ = batch_least_squares(X, y, noise_var=noise_var)
    assert np.allclose(traj[-1], wls, atol=1e-6)


def test_kalman_forgetting_mode_equals_rls_lambda():
    """P <- P/lambda prediction == RLS with forgetting lambda, exactly."""
    rng = np.random.default_rng(13)
    X, y, _, _ = make_problem(rng)
    lam, delta = 0.95, 1e-4
    traj_k, _ = kalman_parameter_filter(X, y, meas_var=1.0,
                                        forgetting=lam, P0=1.0 / delta)
    rls = RecursiveLeastSquares(X.shape[1], forgetting=lam, delta=delta)
    traj_r = rls.fit(X, y)
    assert np.allclose(traj_k, traj_r, atol=1e-8)


def test_kalman_process_noise_keeps_covariance_open():
    """With Q > 0 the posterior variance reaches a floor instead of 1/k decay."""
    rng = np.random.default_rng(14)
    n = 3000
    X = np.ones((n, 1))
    y = rng.normal(size=n)
    _, P_static = kalman_parameter_filter(X, y, meas_var=1.0, P0=1e8)
    _, P_track = kalman_parameter_filter(X, y, meas_var=1.0,
                                         process_var=1e-3, P0=1e8)
    assert P_static[0, 0] == pytest.approx(1.0 / n, rel=0.01)
    # scalar steady state of P + q = predicted, updated with R = 1:
    # p = (sqrt(q^2 + 4q) - q) / 2
    q = 1e-3
    p_ss = (np.sqrt(q**2 + 4 * q) - q) / 2
    assert P_track[0, 0] == pytest.approx(p_ss, rel=0.01)
    assert P_track[0, 0] > 30 * P_static[0, 0]


def test_kalman_tracks_drifting_parameter_where_static_lags():
    rng = np.random.default_rng(15)
    n = 2000
    walk = np.cumsum(rng.normal(0, 0.05, size=n))          # true drifting scalar
    X = np.ones((n, 1))
    y = walk + rng.normal(0, 0.5, size=n)
    traj_s, _ = kalman_parameter_filter(X, y, meas_var=0.25, P0=1e8)
    traj_t, _ = kalman_parameter_filter(X, y, meas_var=0.25,
                                        process_var=0.05**2, P0=1e8)
    err_s = np.mean((traj_s[n // 2:, 0] - walk[n // 2:]) ** 2)
    err_t = np.mean((traj_t[n // 2:, 0] - walk[n // 2:]) ** 2)
    assert err_t < err_s / 5


def test_kalman_validation():
    X = np.ones((5, 1)); y = np.ones(5)
    with pytest.raises(ValueError):
        kalman_parameter_filter(X, y, meas_var=0.0)
    with pytest.raises(ValueError):
        kalman_parameter_filter(X, y, meas_var=1.0, process_var=-1.0)
    with pytest.raises(ValueError):
        kalman_parameter_filter(X, y, meas_var=1.0, forgetting=1.5)


# ----------------------------------------------------- total least squares

def test_ls_attenuation_matches_closed_form():
    """The relationship: LS slope / true slope == kappa, not just 'biased'."""
    rng = np.random.default_rng(16)
    a, sx, su = 2.0, 1.0, 0.6
    n = 200_000
    x_true = rng.normal(0, sx, n)
    x_obs = x_true + rng.normal(0, su, n)
    y = a * x_true + rng.normal(0, 0.1, n)
    a_ls = (x_obs @ y) / (x_obs @ x_obs)
    kappa = attenuation_factor(sx**2, su**2)
    assert a_ls / a == pytest.approx(kappa, rel=0.02)
    assert kappa == pytest.approx(sx**2 / (sx**2 + su**2), rel=1e-12)


def test_tls_recovers_slope_under_equal_eiv_noise():
    rng = np.random.default_rng(17)
    a, sx, s_noise = 2.0, 1.0, 0.6
    n = 200_000
    x_true = rng.normal(0, sx, n)
    A = (x_true + rng.normal(0, s_noise, n))[:, None]
    b = a * x_true + rng.normal(0, s_noise, n)
    x_tls, _ = total_least_squares(A, b)
    assert x_tls[0] == pytest.approx(a, rel=0.02)
    # and LS on the same data shows the predicted attenuation
    a_ls = float(A[:, 0] @ b / (A[:, 0] @ A[:, 0]))
    assert a_ls / a == pytest.approx(attenuation_factor(sx**2, s_noise**2),
                                     rel=0.02)


def test_tls_reduces_to_ls_when_regressors_clean():
    """With noise only in b, both estimators agree (and are consistent)."""
    rng = np.random.default_rng(18)
    X, y, theta, _ = make_problem(rng, n=5000, sigma=0.05)
    x_tls, _ = total_least_squares(X, y)
    ls, _ = batch_least_squares(X, y)
    assert np.allclose(x_tls, ls, atol=0.01)


def test_tls_validation():
    with pytest.raises(ValueError):
        total_least_squares(np.ones((4, 2)), np.ones(3))


# ------------------------------------------------------- FD-TLS criterion

def test_fd_tls_criterion_equals_direct_minimisation():
    """The archived closed form == direct minimisation over the nuisance S."""
    rng = np.random.default_rng(19)
    n = 64
    X1 = rng.normal(size=n) + 1j * rng.normal(size=n)
    X2 = rng.normal(size=n) + 1j * rng.normal(size=n)
    G = 0.7 * (rng.normal(size=n) + 1j * rng.normal(size=n))
    P1 = rng.uniform(0.5, 2.0, n)
    P2 = rng.uniform(0.5, 2.0, n)
    # optimal nuisance spectrum from the derivation (colored case)
    S = (X1 / P1 + X2 * np.conj(G) / P2) / (1 / P1 + np.abs(G)**2 / P2)
    J_direct = np.sum(np.abs(X1 - S)**2 / P1 + np.abs(X2 - S * G)**2 / P2)
    J_closed = fd_tls_criterion(X1, X2, G, P1=P1, P2=P2)
    assert J_closed == pytest.approx(float(J_direct), rel=1e-12)
    # S is a genuine minimiser: perturbing it raises the direct criterion
    S_p = S * 1.001
    J_pert = np.sum(np.abs(X1 - S_p)**2 / P1 + np.abs(X2 - S_p * G)**2 / P2)
    assert J_pert > J_direct


def test_fd_tls_criterion_white_equal_power_special_case():
    rng = np.random.default_rng(20)
    n = 32
    X1 = rng.normal(size=n) + 1j * rng.normal(size=n)
    X2 = rng.normal(size=n) + 1j * rng.normal(size=n)
    G = np.full(n, 0.5 + 0.2j)
    J = fd_tls_criterion(X1, X2, G)
    direct = np.sum(np.abs(X2 - G * X1)**2 / (1 + np.abs(G)**2))
    assert J == pytest.approx(float(direct), rel=1e-12)


def test_fd_tls_finds_true_delay_parameter():
    """End to end: minimising J over a delay model recovers the delay."""
    rng = np.random.default_rng(21)
    n = 256
    f = np.fft.rfftfreq(n, d=1.0)[1:]                   # skip DC
    S = (rng.normal(size=f.size) + 1j * rng.normal(size=f.size)) * 3.0
    d_true = 7.3
    sigma = 0.5
    X1 = S + sigma * (rng.normal(size=f.size) + 1j * rng.normal(size=f.size))
    X2 = S * np.exp(-2j * np.pi * f * d_true) \
        + sigma * (rng.normal(size=f.size) + 1j * rng.normal(size=f.size))
    ds = np.linspace(5.0, 9.0, 2001)
    J = [fd_tls_criterion(X1, X2, np.exp(-2j * np.pi * f * d),
                          P1=sigma**2, P2=sigma**2) for d in ds]
    assert ds[int(np.argmin(J))] == pytest.approx(d_true, abs=0.05)


# --------------------------------------------------------- constrained LS

def test_constrained_ls_satisfies_constraints():
    rng = np.random.default_rng(22)
    X = rng.normal(size=(50, 5))
    y = rng.normal(size=50)
    C = rng.normal(size=(2, 5))
    d = rng.normal(size=2)
    theta = constrained_least_squares(X, y, C, d)
    assert np.allclose(C @ theta, d, atol=1e-9)


def test_constrained_ls_matches_projection_formula():
    """KKT solution == the archived two-inverse Lagrange closed form."""
    rng = np.random.default_rng(23)
    X = rng.normal(size=(40, 4))
    y = rng.normal(size=40)
    C = rng.normal(size=(1, 4))
    d = np.array([2.0])
    theta = constrained_least_squares(X, y, C, d)
    XtX_inv = np.linalg.inv(X.T @ X)
    unc = XtX_inv @ (X.T @ y)
    lam = np.linalg.solve(C @ XtX_inv @ C.T, d - C @ unc)
    closed = XtX_inv @ (X.T @ y + C.T @ lam)
    assert np.allclose(theta, closed, atol=1e-9)


def test_constrained_ls_improves_toward_ls_without_constraints_binding():
    """A constraint already satisfied by the LS solution changes nothing."""
    rng = np.random.default_rng(24)
    X = rng.normal(size=(60, 3))
    y = rng.normal(size=60)
    ls, _ = batch_least_squares(X, y)
    C = rng.normal(size=(1, 3))
    d = C @ ls                                          # non-binding by design
    theta = constrained_least_squares(X, y, C, d)
    assert np.allclose(theta, ls, atol=1e-9)


def test_constrained_ls_polynomial_through_a_point():
    """The archived use case: fit a parabola forced through a known point."""
    rng = np.random.default_rng(25)
    t = np.linspace(0, 1, 40)
    X = np.vander(t, 3, increasing=True)
    y = 1.0 + 2.0 * t - 3.0 * t**2 + 0.1 * rng.normal(size=t.size)
    C = np.array([[1.0, 0.5, 0.25]])                    # value at t = 0.5
    d = np.array([1.6])
    theta = constrained_least_squares(X, y, C, d)
    assert np.polyval(theta[::-1], 0.5) == pytest.approx(1.6, abs=1e-9)