"""Tests for whitened_least_squares solver."""

import numpy as np
import pytest
from numpy.testing import assert_allclose

from gri_utils.geolocation import whitened_least_squares


def test_linear_system_exact_recovery() -> None:
    """Verify exact solution recovery for a linear system."""
    # Linear system: y = A @ x + noise
    # Residual: r(x) = y - A @ x
    # Jacobian: J = -A

    rng = np.random.default_rng(42)

    # True parameters
    x_true = np.array([1.0, 2.0, 3.0])

    # Design matrix (4 measurements, 3 parameters)
    a_matrix = rng.standard_normal((4, 3))

    # Measurements (no noise for exact recovery)
    y = a_matrix @ x_true

    # Unit covariance (unweighted)
    covariance = np.eye(4)

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - a_matrix @ x, -a_matrix

    x0 = np.zeros(3)
    solution, _cov, chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
    )

    assert_allclose(solution, x_true, atol=1e-10)
    # Exact recovery means chi-squared should be ~0
    assert chi_sq < 1e-15


def test_weighted_linear_system() -> None:
    """Verify weighted least squares with non-identity covariance."""
    rng = np.random.default_rng(123)

    x_true = np.array([2.0, -1.0])

    # Design matrix
    a_matrix = np.array(
        [
            [1.0, 0.0],
            [0.0, 1.0],
            [1.0, 1.0],
            [1.0, -1.0],
        ],
    )

    # Measurements with heteroscedastic noise
    sigmas = np.array([0.1, 0.5, 0.2, 0.3])
    noise = rng.standard_normal(4) * sigmas
    y = a_matrix @ x_true + noise

    # Diagonal covariance
    covariance = np.diag(sigmas**2)

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - a_matrix @ x, -a_matrix

    x0 = np.zeros(2)
    solution, sol_cov, _chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
    )

    # Solution should be close to truth (within noise)
    # With 4 measurements of 2 parameters, expect reasonable accuracy
    assert np.linalg.norm(solution - x_true) < 1.0

    # Covariance should be positive definite
    eigvals = np.linalg.eigvalsh(sol_cov)
    assert np.all(eigvals > 0)


def test_nonlinear_system() -> None:
    """Verify convergence for a simple nonlinear system."""
    # y = x^2 (scalar case wrapped in arrays)
    x_true = np.array([3.0])
    y = np.array([9.0])  # 3^2 = 9
    covariance = np.array([[0.01]])

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - x**2, -2 * x.reshape(1, 1)

    # Start from reasonable initial guess
    x0 = np.array([2.0])
    solution, _cov, _chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
    )

    assert_allclose(solution, x_true, atol=1e-6)


def test_covariance_output_shape() -> None:
    """Verify solution covariance has correct shape."""
    # Simple 2-parameter, 3-measurement system
    a_matrix = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
    y = np.array([1.0, 2.0, 3.0])
    covariance = np.eye(3)

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - a_matrix @ x, -a_matrix

    x0 = np.zeros(2)
    solution, sol_cov, chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
    )

    assert solution.shape == (2,)
    assert sol_cov.shape == (2, 2)
    assert isinstance(chi_sq, float)


def test_correlated_measurements() -> None:
    """Verify handling of correlated measurement covariance."""
    x_true = np.array([1.0, 2.0])

    a_matrix = np.array([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]])
    y = a_matrix @ x_true

    # Correlated covariance matrix
    covariance = np.array(
        [
            [1.0, 0.3, 0.1],
            [0.3, 1.0, 0.3],
            [0.1, 0.3, 1.0],
        ],
    )

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - a_matrix @ x, -a_matrix

    x0 = np.zeros(2)
    solution, sol_cov, _chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
    )

    assert_allclose(solution, x_true, atol=1e-10)
    assert sol_cov.shape == (2, 2)


def test_non_positive_definite_covariance_raises() -> None:
    """Verify error on non-positive-definite covariance."""
    # Singular covariance matrix
    covariance = np.array([[1.0, 1.0], [1.0, 1.0]])

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return np.array([1.0, 2.0]) - x, -np.eye(2)

    x0 = np.zeros(2)

    with pytest.raises(np.linalg.LinAlgError):
        whitened_least_squares(residual_and_jacobian, x0, covariance)


def test_least_squares_kwargs_passthrough() -> None:
    """Verify kwargs are passed to scipy.optimize.least_squares."""
    a_matrix = np.array([[1.0, 0.0], [0.0, 1.0]])
    y = np.array([1.0, 2.0])
    covariance = np.eye(2)

    def residual_and_jacobian(
        x: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        return y - a_matrix @ x, -a_matrix

    x0 = np.zeros(2)

    # Pass custom tolerances
    solution, _cov, _chi_sq = whitened_least_squares(
        residual_and_jacobian,
        x0,
        covariance,
        ftol=1e-12,
        xtol=1e-12,
    )

    assert_allclose(solution, y, atol=1e-10)
