"""Predicted solution covariance from geometry and measurement uncertainty.

Computes the theoretical covariance of a weighted least squares solution
given the measurement Jacobian and measurement covariance matrix. This is
the analytical equivalent of running a full solver -- it predicts how
measurement uncertainty propagates through geometry into solution uncertainty.

This is closely related to Geometric Dilution of Precision (GDOP): the
Jacobian encodes the geometry, and the measurement covariance encodes the
sensor noise model. The output tells you how well a given sensor/geometry
configuration can resolve the unknowns.

Typical usage:

    1. Evaluate a Jacobian at the candidate geometry using any *_with_gradient
       function (tdoa_with_gradient, aoa_quat_with_gradient, etc.)
    2. Build a measurement covariance matrix (diagonal for independent
       measurements, full matrix for correlated measurements)
    3. Call predicted_covariance to get the (N, N) solution covariance

The result can be used to:
    - Compare sensor configurations without running Monte Carlo
    - Compute GDOP, PDOP, HDOP, VDOP from the solution covariance
    - Evaluate whether a system design meets accuracy requirements
    - Perform sensitivity analysis by sweeping geometry or noise parameters
"""

from __future__ import annotations

from typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    from numpy.typing import NDArray


def predicted_covariance(
    jacobian: NDArray[np.floating],
    measurement_covariance: NDArray[np.floating],
) -> NDArray[np.floating]:
    """Compute predicted solution covariance from Jacobian and measurement noise.

    Given a measurement Jacobian J (M x N) and measurement covariance R (M x M),
    computes the solution covariance:

        C = (J^T R^{-1} J)^{-1}

    This is the Cramer-Rao lower bound for an unbiased estimator -- the best
    achievable covariance for any weighted least squares solution with this
    geometry and measurement noise.

    The computation uses Cholesky whitening for numerical stability:
        L = cholesky(R)
        J_w = L^{-1} J
        C = (J_w^T J_w)^{-1}

    Args:
        jacobian: Measurement Jacobian matrix. Shape (M, N) where M is the
            number of measurements and N is the number of unknowns (typically 3
            for position or 6 for position+velocity). Each row is the gradient
            of one measurement with respect to the unknowns.
        measurement_covariance: Measurement covariance matrix. Shape (M, M).
            Must be symmetric positive definite. Can be diagonal (independent
            measurements) or full (correlated measurements).

    Returns:
        Predicted solution covariance matrix. Shape (N, N). For position-only
        problems this is a 3x3 ECEF covariance.

    Raises:
        LinAlgError: If measurement_covariance is not positive definite.
        LinAlgError: If the Jacobian is rank-deficient (geometry cannot
            resolve all unknowns).

    Example:
        Predict TDOA geolocation accuracy for a given geometry::

            from gri_utils.observables import tdoa_with_gradient
            from gri_utils.geolocation import predicted_covariance
            import numpy as np

            # Candidate target position and 3 collector pairs
            target = np.array([6378137.0, 100.0, 100.0])
            c1 = np.array([[6378137.0, 1000.0, 0.0],
                           [6378137.0, 0.0, 1000.0],
                           [6378137.0, -1000.0, 0.0]])
            c2 = np.array([[6378137.0, -1000.0, 0.0],
                           [6378137.0, 0.0, -1000.0],
                           [6378137.0, 1000.0, 0.0]])

            # Get Jacobian at this geometry
            jacobians = np.empty((3, 3))
            for i in range(3):
                _, jacobians[i] = tdoa_with_gradient(target, c1[i], c2[i])

            # 10 ns timing uncertainty, independent measurements
            R = np.diag(np.full(3, (10e-9)**2))

            # Predicted position covariance (3x3, ECEF)
            cov = predicted_covariance(jacobians, R)
            print(f"Position stdev: {np.sqrt(np.diag(cov))} meters")
    """
    jacobian = np.asarray(jacobian)
    measurement_covariance = np.asarray(measurement_covariance)

    # Whiten the Jacobian via Cholesky decomposition
    chol_lower = cholesky(measurement_covariance, lower=True)
    j_white = solve_triangular(chol_lower, jacobian, lower=True)

    # Solution covariance = (J_w^T J_w)^{-1}
    jtj = j_white.T @ j_white
    return np.linalg.inv(jtj)
