"""Recursive estimation: batch least squares, RLS, and the Kalman connection.

The tools behind the recursive-estimation topic page: the batch and ridge
least-squares baselines, the classical RLS recursion (sample-by-sample and
in rank-B blocks, following the author's MSc-thesis code), the Kalman filter
restricted to a parameter state (which reproduces recursive weighted least
squares exactly), and the total-least-squares family for the case where the
regressors themselves are noisy.

Two identities anchor everything and are pinned by the test suite:

* RLS with forgetting factor 1 and a vanishing initial regularisation equals
  batch least squares; with a finite regularisation ``delta`` it equals ridge
  regression with ridge weight ``delta``.  Both are exact identities; the
  tests pin them at 1e-7 and 1e-12 respectively.
* The Kalman filter on a static parameter (F = I, Q = 0) with per-sample
  measurement variances equals recursive weighted least squares, and hence
  batch WLS (exact identity; pinned at 1e-6 by the tests).

Convention note (the arc's parameterisation trap, RLS edition): this module
initialises the inverse-correlation matrix as ``P0 = I / delta`` with
``delta`` SMALL, the convention of the thesis code (``rlsinit.m``,
``MISO_brls.m``).  Some texts write the same thing as ``P0 = delta * I``
with delta LARGE.  The quantity ``delta`` here is exactly the ridge weight
added to X^T X, which is what makes the first identity exact.
"""

import numpy as np


def batch_least_squares(X: np.ndarray, y: np.ndarray,
                        noise_var: np.ndarray | float | None = None,
                        ridge: float = 0.0) -> tuple[np.ndarray, np.ndarray]:
    """Solve ``y = X @ theta + n`` in one batch, optionally weighted or ridged.

    Parameters
    ----------
    X : ndarray of shape (n_samples, n_params)
        Design matrix (one regressor row per observation).
    y : ndarray of shape (n_samples,)
        Observations.
    noise_var : ndarray of shape (n_samples,), float, or None
        Known noise variance per sample; weights are their inverses.
        None or a scalar reduces to ordinary least squares.
    ridge : float
        Non-negative weight of an added ``ridge * I`` term in the normal
        equations.  ``ridge > 0`` is exactly what a finite RLS
        initialisation ``P0 = I / ridge`` computes.

    Returns
    -------
    theta : ndarray of shape (n_params,)
        The estimate ``(X^T W X + ridge I)^-1 X^T W y``.
    P : ndarray of shape (n_params, n_params)
        The matrix ``(X^T W X + ridge I)^-1``.  For ``noise_var`` equal to
        the true per-sample variances and ``ridge = 0`` this is the exact
        covariance of theta and equals the Cramer-Rao bound.

    Raises
    ------
    ValueError
        If shapes disagree, a variance is non-positive, ridge is negative,
        or the problem is underdetermined with ``ridge = 0``.
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    n_samples, n_params = X.shape
    if y.size != n_samples:
        raise ValueError(f"X has {n_samples} rows but y has {y.size} entries")
    if ridge < 0:
        raise ValueError("ridge must be non-negative")
    if n_samples < n_params and ridge == 0.0:
        raise ValueError("underdetermined: fewer observations than parameters")

    if noise_var is None:
        w = np.ones(n_samples)
    else:
        var = np.broadcast_to(np.asarray(noise_var, dtype=float), (n_samples,))
        if np.any(var <= 0):
            raise ValueError("noise_var must be strictly positive")
        w = 1.0 / var

    XtW = X.T * w
    P = np.linalg.inv(XtW @ X + ridge * np.eye(n_params))
    theta = P @ (XtW @ y)
    return theta, P


class RecursiveLeastSquares:
    """Classical RLS: the batch least-squares solution, updated per sample.

    Maintains the estimate ``theta`` and the matrix ``P``, which at
    forgetting factor 1 is ``(delta I + sum x x^T)^-1``: the inverse of the
    (ridge-regularised) normal-equation matrix, so that ``sigma^2 P`` is
    the estimate's covariance once ``delta`` is negligible.

    Parameters
    ----------
    n_params : int
        Number of parameters (length of the regressor vector).
    forgetting : float
        Forgetting factor lambda in (0, 1].  1 keeps all data (growing
        window, consistent estimator); smaller values weight sample ``i``
        by ``lambda^(k-i)``, trading a variance floor for tracking.
    delta : float
        Initial regularisation: ``P0 = I / delta`` (thesis convention,
        delta SMALL).  At lambda = 1 the recursion then computes ridge
        regression with ridge weight delta, exactly.
    theta0 : ndarray of shape (n_params,), optional
        Initial estimate (default zeros).

    Notes
    -----
    The per-sample update, with a priori error ``e = y - x^T theta``:

    .. code-block:: text

        g = P x / (lambda + x^T P x)
        theta <- theta + g e
        P <- (P - g x^T P) / lambda

    This is the Sherman-Morrison rank-one update of the inverse
    normal-equation matrix; no matrix is ever inverted at run time.
    """

    def __init__(self, n_params: int, forgetting: float = 1.0,
                 delta: float = 1e-6, theta0: np.ndarray | None = None):
        if not 0.0 < forgetting <= 1.0:
            raise ValueError("forgetting factor must be in (0, 1]")
        if delta <= 0:
            raise ValueError("delta must be positive")
        self.lam = float(forgetting)
        self.theta = (np.zeros(n_params) if theta0 is None
                      else np.asarray(theta0, dtype=float).copy())
        if self.theta.shape != (n_params,):
            raise ValueError("theta0 must have shape (n_params,)")
        self.P = np.eye(n_params) / delta

    def update(self, x: np.ndarray, y: float) -> float:
        """Fold in one observation; return the a priori error.

        Parameters
        ----------
        x : ndarray of shape (n_params,)
            Regressor vector for this observation.
        y : float
            Observed response.

        Returns
        -------
        e : float
            A priori error ``y - x^T theta`` (before the update), the
            quantity RLS applications call the innovation or residual.
        """
        x = np.asarray(x, dtype=float).ravel()
        Px = self.P @ x
        g = Px / (self.lam + x @ Px)
        e = float(y - x @ self.theta)
        self.theta = self.theta + g * e
        self.P = (self.P - np.outer(g, x @ self.P)) / self.lam
        return e

    def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Run the recursion over a whole record; return the trajectory.

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_params)
        y : ndarray of shape (n_samples,)

        Returns
        -------
        trajectory : ndarray of shape (n_samples, n_params)
            ``trajectory[k]`` is the estimate after observation k.
        """
        X = np.atleast_2d(np.asarray(X, dtype=float))
        y = np.asarray(y, dtype=float).ravel()
        if X.shape[0] != y.size:
            raise ValueError("X and y disagree on the number of samples")
        traj = np.empty_like(X)
        for k in range(X.shape[0]):
            self.update(X[k], y[k])
            traj[k] = self.theta
        return traj


def block_rls(X: np.ndarray, y: np.ndarray, block_size: int,
              delta: float = 1e-8) -> tuple[np.ndarray, np.ndarray]:
    """Block RLS: fold in B observations per step (rank-B Woodbury update).

    The block generalisation of the Sherman-Morrison step used by the
    author's MSc code (``MISO_brls.m``): per block, invert (implicitly, via
    ``solve``) a B x B matrix instead of running B rank-one updates.  No
    forgetting: this is the growing-window algorithm, and its fixed point
    is the batch (left pseudo-inverse) solution.

    Parameters
    ----------
    X : ndarray of shape (n_samples, n_params)
    y : ndarray of shape (n_samples,)
    block_size : int
        Observations per update (B = 1 reproduces sample RLS exactly).
        Trailing samples that do not fill a block are ignored, matching
        the source's loop structure.
    delta : float
        Initial regularisation, ``P0 = I / delta``.

    Returns
    -------
    theta : ndarray of shape (n_params,)
    P : ndarray of shape (n_params, n_params)

    Notes
    -----
    Numerical caveat, quantified on the topic page: when ``block_size``
    exceeds ``n_params`` the B x B matrix ``I + Xb P Xb^T`` has
    ``n_params`` eigenvalues of order ``x^T x / delta`` and the rest equal
    to 1, so its condition number is of order ``1/delta``.  Solving the
    linear system (as here) keeps the error near machine precision;
    explicitly inverting it (as the 2001 MATLAB did) loses about
    ``log10(1/delta)`` digits.
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    n_samples, n_params = X.shape
    if X.shape[0] != y.size:
        raise ValueError("X and y disagree on the number of samples")
    if not 1 <= block_size <= n_samples:
        raise ValueError("block_size must be in [1, n_samples]")
    if delta <= 0:
        raise ValueError("delta must be positive")

    theta = np.zeros(n_params)
    P = np.eye(n_params) / delta
    B = block_size
    for k in range(0, n_samples - B + 1, B):
        Xb = X[k:k + B]
        yb = y[k:k + B]
        PXt = P @ Xb.T                                   # (p, B)
        M = np.eye(B) + Xb @ PXt                         # I + Xb P Xb^T
        G = np.linalg.solve(M.T, PXt.T).T                # P Xb^T M^-1
        theta = theta + G @ (yb - Xb @ theta)
        P = P - G @ (Xb @ P)
    return theta, P


def kalman_parameter_filter(X: np.ndarray, y: np.ndarray,
                            meas_var: np.ndarray | float,
                            process_var: float = 0.0,
                            forgetting: float | None = None,
                            theta0: np.ndarray | None = None,
                            P0: np.ndarray | float = 1e8
                            ) -> tuple[np.ndarray, np.ndarray]:
    """Kalman filter whose state is a parameter vector: recursive WLS and beyond.

    The state-space model is ``theta_k = theta_{k-1} + w_k`` (random walk,
    F = I, Q = process_var * I) observed through ``y_k = x_k^T theta_k + v_k``
    with ``var(v_k) = meas_var[k]``.  Three regimes:

    * ``process_var = 0``: static parameter.  The filter IS recursive
      weighted least squares and matches batch WLS (exact identity).
    * ``process_var > 0``: the parameter drifts; the covariance no longer
      shrinks to zero and the filter keeps tracking.
    * ``forgetting`` set (overrides process_var): the prediction step
      divides P by lambda, which reproduces RLS with that forgetting
      factor exactly; equivalently a state-proportional
      ``Q = (1/lambda - 1) P``.

    Parameters
    ----------
    X : ndarray of shape (n_samples, n_params)
        Measurement rows x_k^T.
    y : ndarray of shape (n_samples,)
    meas_var : ndarray of shape (n_samples,) or float
        Measurement noise variance per sample (scalar broadcasts).
    process_var : float
        Random-walk increment variance q (per parameter, isotropic).
    forgetting : float or None
        If set, use the RLS-equivalent prediction ``P <- P / lambda``.
    theta0 : ndarray of shape (n_params,), optional
        Initial state estimate (default zeros).
    P0 : ndarray of shape (n_params, n_params) or float
        Initial covariance (scalar means ``P0 * I``).  Large P0 is the
        diffuse prior; it plays exactly the role of ``1/delta`` in RLS.

    Returns
    -------
    trajectory : ndarray of shape (n_samples, n_params)
        State estimate after each measurement.
    P : ndarray of shape (n_params, n_params)
        Final posterior covariance.
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    n_samples, n_params = X.shape
    if X.shape[0] != y.size:
        raise ValueError("X and y disagree on the number of samples")
    r = np.broadcast_to(np.asarray(meas_var, dtype=float), (n_samples,))
    if np.any(r <= 0):
        raise ValueError("meas_var must be strictly positive")
    if process_var < 0:
        raise ValueError("process_var must be non-negative")
    if forgetting is not None and not 0.0 < forgetting <= 1.0:
        raise ValueError("forgetting factor must be in (0, 1]")

    theta = (np.zeros(n_params) if theta0 is None
             else np.asarray(theta0, dtype=float).copy())
    P = (np.eye(n_params) * float(P0) if np.isscalar(P0)
         else np.asarray(P0, dtype=float).copy())
    traj = np.empty((n_samples, n_params))
    for k in range(n_samples):
        # predict: F = I, so only the covariance moves
        if forgetting is not None:
            P = P / forgetting
        elif process_var > 0.0:
            P = P + process_var * np.eye(n_params)
        # update
        x = X[k]
        Px = P @ x
        gain = Px / (x @ Px + r[k])
        theta = theta + gain * (y[k] - x @ theta)
        P = P - np.outer(gain, x @ P)
        traj[k] = theta
    return traj, P


def ew_variance_factor(forgetting: float) -> float:
    """Steady-state variance of the exponentially weighted mean, over sigma^2.

    Estimating a constant from unit-variance noise with forgetting factor
    lambda gives asymptotic variance ``sigma^2 (1 - lambda) / (1 + lambda)``:
    the geometric weights ``(1-lambda) lambda^j`` have sum-of-squares
    ``(1-lambda)/(1+lambda)``.  This is the price of tracking: at lambda = 1
    the variance falls as 1/k forever; at lambda < 1 it stops falling here.

    Parameters
    ----------
    forgetting : float
        Forgetting factor lambda in (0, 1).

    Returns
    -------
    factor : float
        ``(1 - lambda) / (1 + lambda)``, to be multiplied by the
        measurement noise variance.
    """
    if not 0.0 < forgetting < 1.0:
        raise ValueError("forgetting factor must be in (0, 1) for a "
                         "finite-memory variance")
    return (1.0 - forgetting) / (1.0 + forgetting)


def variance_equivalent_window(forgetting: float) -> float:
    """Number of equally weighted samples with the same variance as EW averaging.

    Two "effective window" conventions coexist for a forgetting factor
    lambda, and they differ by a factor of about two:

    * weight-sum window: ``sum lambda^j = 1 / (1 - lambda)``,
    * variance-equivalent window: the N for which a plain mean of N
      samples has the same variance, ``(1 + lambda) / (1 - lambda)``,
      approximately ``2 / (1 - lambda)``.

    Quoting a bound with the wrong one silently costs a factor of two,
    the same trap as the SNR conventions in the sinusoid CRLB.  This
    function returns the variance-equivalent one.

    Parameters
    ----------
    forgetting : float
        Forgetting factor lambda in (0, 1).

    Returns
    -------
    n_eff : float
        ``(1 + lambda) / (1 - lambda)``.
    """
    if not 0.0 < forgetting < 1.0:
        raise ValueError("forgetting factor must be in (0, 1)")
    return (1.0 + forgetting) / (1.0 - forgetting)


def total_least_squares(A: np.ndarray, b: np.ndarray
                        ) -> tuple[np.ndarray, float]:
    """Total least squares solution of ``A x ~ b`` via the SVD.

    Ordinary LS perturbs only b to make the system consistent; TLS finds
    the smallest joint perturbation of [A b] (in Frobenius norm), which is
    the right criterion when the regressors are as noisy as the
    observations (errors-in-variables, orthogonal regression).

    Parameters
    ----------
    A : ndarray of shape (n_samples, n_params)
        Noisy design matrix.
    b : ndarray of shape (n_samples,)
        Noisy observations.  Both are assumed to carry iid noise of the
        SAME variance; with unequal variances, scale the columns first
        (or the estimate remains biased).

    Returns
    -------
    x : ndarray of shape (n_params,)
        TLS estimate: from the right singular vector of [A b] belonging
        to the smallest singular value.
    sigma_min : float
        That smallest singular value; its square is the residual
        (the size of the minimal perturbation).

    Raises
    ------
    ValueError
        On shape mismatch, or when the smallest right singular vector has
        (numerically) zero weight on b, in which case the TLS solution
        does not exist in this generic form.
    """
    A = np.atleast_2d(np.asarray(A, dtype=float))
    b = np.asarray(b, dtype=float).ravel()
    if A.shape[0] != b.size:
        raise ValueError("A and b disagree on the number of samples")
    Z = np.column_stack([A, b])
    _, s, Vt = np.linalg.svd(Z, full_matrices=False)
    v = Vt[-1]
    if abs(v[-1]) < 1e-12:
        raise ValueError("degenerate TLS problem: smallest singular vector "
                         "has no component on b")
    return -v[:-1] / v[-1], float(s[-1])


def attenuation_factor(signal_var: float, error_var: float) -> float:
    """The multiplicative bias of LS when the regressor itself is noisy.

    For ``y = a x_true + n`` observed through ``x = x_true + u`` (all
    zero-mean, independent), the LS slope converges not to ``a`` but to
    ``a * signal_var / (signal_var + error_var)``: regression dilution.
    More data does not help; the estimator is inconsistent.

    Parameters
    ----------
    signal_var : float
        Variance of the true (noise-free) regressor.
    error_var : float
        Variance of the regressor's measurement error.

    Returns
    -------
    kappa : float
        ``signal_var / (signal_var + error_var)`` in (0, 1].
    """
    if signal_var <= 0 or error_var < 0:
        raise ValueError("signal_var must be positive, error_var non-negative")
    return signal_var / (signal_var + error_var)


def constrained_least_squares(X: np.ndarray, y: np.ndarray,
                              C: np.ndarray, d: np.ndarray) -> np.ndarray:
    """Least squares subject to exact linear constraints ``C theta = d``.

    Solves ``min ||y - X theta||^2  s.t.  C theta = d`` via the KKT
    system; the classic use is polynomial fitting pinned to pass through
    known points (the archived derivation this follows).

    Parameters
    ----------
    X : ndarray of shape (n_samples, n_params)
    y : ndarray of shape (n_samples,)
    C : ndarray of shape (n_constraints, n_params)
        Constraint matrix; rows must be linearly independent.
    d : ndarray of shape (n_constraints,)
        Constraint values.

    Returns
    -------
    theta : ndarray of shape (n_params,)
        The constrained minimiser; satisfies ``C theta = d`` to solver
        precision.
    """
    X = np.atleast_2d(np.asarray(X, dtype=float))
    y = np.asarray(y, dtype=float).ravel()
    C = np.atleast_2d(np.asarray(C, dtype=float))
    d = np.asarray(d, dtype=float).ravel()
    n_samples, n_params = X.shape
    n_con = C.shape[0]
    if y.size != n_samples:
        raise ValueError("X and y disagree on the number of samples")
    if C.shape[1] != n_params or d.size != n_con:
        raise ValueError("constraint shapes disagree with n_params")

    kkt = np.block([[2.0 * X.T @ X, C.T],
                    [C, np.zeros((n_con, n_con))]])
    rhs = np.concatenate([2.0 * X.T @ y, d])
    return np.linalg.solve(kkt, rhs)[:n_params]


def fd_tls_criterion(X1: np.ndarray, X2: np.ndarray, G: np.ndarray,
                     P1: np.ndarray | float = 1.0,
                     P2: np.ndarray | float = 1.0) -> float:
    """Concentrated errors-in-variables criterion for a two-channel model.

    For the frequency-domain model ``X1 = S + N1``, ``X2 = S G + N2`` with
    per-bin noise powers P1 and P2, minimising the (inverse-variance
    weighted) noise energy over the nuisance spectrum S leaves

        J(G) = sum |X2 - G X1|^2 / (P2 + P1 |G|^2),

    the archived TLS derivation's closed form (re-verified for this page
    against direct minimisation over S).  LS against channel 1 would use
    denominator P2 alone; the extra ``P1 |G|^2`` is what stops the noisy
    reference channel from biasing the fit.

    Parameters
    ----------
    X1, X2 : ndarray of shape (n_bins,), complex
        Measured channel spectra.
    G : ndarray of shape (n_bins,), complex
        Candidate transfer function (e.g. ``exp(-gamma(f) d)`` for a
        propagation model with parameter d).
    P1, P2 : ndarray of shape (n_bins,) or float
        Noise power spectra of the two channels (scalars broadcast).

    Returns
    -------
    J : float
        The criterion value; minimise over the parameters of G.
    """
    X1 = np.asarray(X1); X2 = np.asarray(X2); G = np.asarray(G)
    if not X1.shape == X2.shape == G.shape:
        raise ValueError("X1, X2, G must share one shape")
    num = np.abs(X2 - G * X1) ** 2
    den = np.broadcast_to(P2, X1.shape) + np.broadcast_to(P1, X1.shape) * np.abs(G) ** 2
    if np.any(den <= 0):
        raise ValueError("noise powers must be strictly positive")
    return float(np.sum(num / den))