"""Linear-algebra helper functions for the Ellipsoidal Toolbox.

Ports of MATLAB utilities: ``ell_valign``, ``ell_simdiag``,
``ell_inv``, and ``regularize``.
"""

from __future__ import annotations

import numpy as np
import scipy.linalg

from .config import ABS_TOL

__all__ = [
    "regularize",
    "ell_valign",
    "ell_simdiag",
    "ell_inv",
]


def regularize(Q: np.ndarray, delta: float = ABS_TOL) -> np.ndarray:
    """Regularize a singular symmetric matrix by adding *delta* * I.

    Parameters
    ----------
    Q : ndarray (n, n)
        Symmetric matrix (possibly rank-deficient).
    delta : float
        Regularization parameter.

    Returns
    -------
    ndarray (n, n)
        Full-rank matrix ``Q + delta * I``.
    """
    n = Q.shape[0]
    return Q + delta * np.eye(n)


def ell_inv(Q: np.ndarray) -> np.ndarray:
    """Safe matrix inverse with regularization fallback.

    If *Q* is singular, it is regularized before inversion.

    Parameters
    ----------
    Q : ndarray (n, n)

    Returns
    -------
    ndarray (n, n)
        Inverse (or pseudo-inverse via regularization).
    """
    if np.linalg.matrix_rank(Q, tol=ABS_TOL) < Q.shape[0]:
        Q = regularize(Q)
    return np.linalg.inv(Q)


def ell_valign(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
    """Compute orthogonal matrix *S* such that S v1 is parallel to v2.

    Used in Minkowski internal approximations where vectors from
    different ellipsoids must be aligned.

    Parameters
    ----------
    v1, v2 : ndarray (n,)
        Input vectors.

    Returns
    -------
    ndarray (n, n)
        Orthogonal matrix *S*.
    """
    n = v1.shape[0]
    n1 = np.linalg.norm(v1)
    n2 = np.linalg.norm(v2)

    if n1 < ABS_TOL or n2 < ABS_TOL:
        return np.eye(n)

    u1 = v1 / n1
    u2 = v2 / n2

    # S = I + (u2 - u1) u1^T  + correction for orthogonality
    # Use the SVD approach: S = V U^T from SVD of u2 u1^T
    M = np.outer(u2, u1)
    U, _, Vt = np.linalg.svd(M)
    S = U @ Vt

    # Ensure proper rotation (det = +1)
    if np.linalg.det(S) < 0:
        U[:, -1] *= -1
        S = U @ Vt

    return S


def ell_simdiag(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """Simultaneous diagonalization transform for two PSD matrices.

    Finds *T* such that ``T A T' = I`` and ``T B T'`` is diagonal.
    Used in ``is_bigger`` to compare ellipsoid sizes.

    Parameters
    ----------
    A : ndarray (n, n)
        Symmetric positive definite matrix (the "bigger" shape).
    B : ndarray (n, n)
        Symmetric positive semi-definite matrix.

    Returns
    -------
    ndarray (n, n)
        Transformation matrix *T*.
    """
    # Generalized eigendecomposition: B v = lambda A v
    # scipy.linalg.eigh(B, A) returns eigenvalues and eigenvectors
    # such that V^T A V = I and V^T B V = diag(eigenvalues).
    # Then T = V^{-T} = (V^T)^{-1}, but since V^T A V = I,
    # we have T = V^T with A-orthonormality.
    #
    # Actually we want T such that T A T' = I, so T = D^{-1/2} U^T
    # from eigendecomposition A = U D U^T.
    eigvals_A, U_A = np.linalg.eigh(A)
    eigvals_A = np.maximum(eigvals_A, ABS_TOL)
    D_inv_sqrt = np.diag(1.0 / np.sqrt(eigvals_A))
    W = D_inv_sqrt @ U_A.T  # W A W^T = I

    B_tilde = W @ B @ W.T
    B_tilde = 0.5 * (B_tilde + B_tilde.T)

    _, U_B = np.linalg.eigh(B_tilde)
    T = U_B.T @ W  # T A T^T = I, T B T^T = diagonal

    return T
