"""Generic whitened nonlinear least squares solver.

Wraps scipy.optimize.least_squares with Cholesky whitening for weighted
least squares problems with correlated measurements.

Returns solution, covariance, and chi-squared fit statistic.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import numpy as np
from scipy.linalg import cholesky, solve_triangular
from scipy.optimize import least_squares

if TYPE_CHECKING:
    from collections.abc import Callable

    from numpy.typing import NDArray


def whitened_least_squares(
    residual_and_jacobian_fn: Callable[
        [NDArray[np.floating]],
        tuple[NDArray[np.floating], NDArray[np.floating]],
    ],
    x0: NDArray[np.floating],
    covariance: NDArray[np.floating],
    **least_squares_kwargs: Any,  # noqa: ANN401
) -> tuple[NDArray[np.floating], NDArray[np.floating], float]:
    """Solve weighted nonlinear least squares with Cholesky whitening.

    Minimizes the weighted sum of squared residuals:
        chi^2(x) = r(x).T @ R^{-1} @ r(x)

    where r(x) is the residual vector and R is the measurement covariance.

    The problem is transformed to standard least squares by whitening:
        r_white = L^{-1} @ r(x)
        J_white = L^{-1} @ J(x)

    where L is the lower Cholesky factor of R (R = L @ L.T).

    Args:
        residual_and_jacobian_fn: Function that computes both residual vector
            r(x) and Jacobian J(x). Returns tuple (residual, jacobian) with
            shapes (M,) and (M, N) respectively.
        x0: Initial guess for solution. Shape (N,).
        covariance: Measurement covariance matrix. Shape (M, M).
        **least_squares_kwargs: Additional arguments passed to
            scipy.optimize.least_squares (e.g., bounds, method, ftol, xtol).

    Returns:
        Tuple of (solution, solution_covariance, chi_squared):
            - solution: Estimated parameter vector. Shape (N,).
            - solution_covariance: Estimated covariance of solution. Shape (N, N).
            - chi_squared: Sum of squared whitened residuals at solution.
              Under correct model assumptions, follows chi-squared distribution
              with (M - N) degrees of freedom.

    Raises:
        LinAlgError: If covariance matrix is not positive definite.
        ValueError: If least_squares fails to converge.

    Note:
        The solution covariance is computed as (J_white.T @ J_white)^{-1}
        evaluated at the solution point.
    """
    x0 = np.asarray(x0)
    covariance = np.asarray(covariance)

    # Compute Cholesky factor once (reused for all iterations)
    chol_lower = cholesky(covariance, lower=True)

    # Cache to avoid duplicate computation when scipy calls residual then jacobian
    cache: dict[str, Any] = {"x": None, "residual": None, "jacobian": None}

    def _cached_residual_and_jacobian(x: NDArray[np.floating]) -> None:
        """Compute residual and jacobian if x changed, caching results."""
        if cache["x"] is None or not np.array_equal(x, cache["x"]):
            residual, jacobian = residual_and_jacobian_fn(x)
            cache["x"] = x.copy()
            cache["residual"] = solve_triangular(chol_lower, residual, lower=True)
            cache["jacobian"] = solve_triangular(chol_lower, jacobian, lower=True)

    def whitened_residual(x: NDArray[np.floating]) -> NDArray[np.floating]:
        _cached_residual_and_jacobian(x)
        return cache["residual"]

    def whitened_jacobian(x: NDArray[np.floating]) -> NDArray[np.floating]:
        _cached_residual_and_jacobian(x)
        return cache["jacobian"]

    # Set default method if not specified
    if "method" not in least_squares_kwargs:
        least_squares_kwargs["method"] = "lm"

    # Run optimizer
    result = least_squares(
        whitened_residual,
        x0,
        jac=whitened_jacobian,  # type: ignore[arg-type]  # scipy stubs incomplete
        **least_squares_kwargs,
    )

    if not result.success:
        msg = f"least_squares failed to converge: {result.message}"
        raise ValueError(msg)

    solution = result.x

    # Compute solution covariance: (J_white.T @ J_white)^{-1}
    # Use pseudoinverse for rank-deficient systems (e.g., FDOA position)
    # This handles cases where the Jacobian doesn't have full column rank
    j_white = whitened_jacobian(solution)
    jtj = j_white.T @ j_white

    # Use pseudoinverse with explicit rcond for numerical stability
    # This sets singular values < rcond * max_sv to zero before inverting
    # A larger rcond (1e-6) helps with ill-conditioned problems like FDOA
    solution_covariance = np.linalg.pinv(jtj, rcond=1e-6)

    # Compute chi-squared: sum of squared whitened residuals
    r_white = whitened_residual(solution)
    chi_squared = float(r_white @ r_white)

    return solution, solution_covariance, chi_squared
