"""Helper to get sorted, normalized eigenvalues and unit vectors from matrices."""

from __future__ import annotations

from typing import TYPE_CHECKING, NamedTuple

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


class Decomposition(NamedTuple):
    """Result of eigenvalue decomposition.

    Attributes:
        vals: Eigenvalues sorted by magnitude (largest first). Shape (n,) for single
            matrix, (..., n) for batched.
        uvecs: Unit eigenvectors as rows, sorted to match vals. Shape (n, n) for single
            matrix, (..., n, n) for batched.
    """

    vals: NDArray[np.floating]
    uvecs: NDArray[np.floating]


def decompose(
    mat: NDArray[np.floating] | Sequence[Sequence[float]],
) -> Decomposition:
    """Return the Eigen tuples (vals,uvecs) in magnitude (greatest to smallest) order.

    Supports single matrices of shape (n, n) or batched matrices of shape (..., n, n).
    Assumes symmetric matrices (uses eigh for efficiency).

    The eigenval is the square of the magnitude. To get the eigenvectors, just use:
        vals, uvecs = decompose(mat)
        vecs = uvecs * (vals**0.5)[..., None]  # Scale unit vectors to eigenvectors
        sma = uvecs[0] * vals[0]**0.5  # First (largest) eigenvector

    Note: Eigenvectors are normalized to point "up" (positive in final axis component).

    Args:
        mat: An array-like structure in (..., n, n) symmetric form. Supports 2x2, 3x3,
            or batched arrays like (N, 2, 2) or (N, 3, 3).

    Returns:
        Decomposition named tuple with:
            - vals: Eigenvalues sorted by magnitude (largest first)
            - uvecs: Unit eigenvectors as rows, sorted to match vals
    """
    mat = np.asarray(mat)

    # Track if single matrix for return shape
    single = mat.ndim == 2  # noqa: PLR2004

    # Ensure at least 3D for consistent processing
    if single:
        mat = mat[np.newaxis, ...]

    # Use eigh (assumes symmetric - faster and more stable)
    vals, uvecs = np.linalg.eigh(mat)  # (..., n), (..., n, n)
    uvecs = np.swapaxes(uvecs, -1, -2)  # Change to row-ordered: (..., n, n)

    # Convert negative eigenvalues to positive (flip eigenvector)
    neg_mask = vals < 0
    vals = np.where(neg_mask, -vals, vals)
    uvecs = np.where(neg_mask[..., np.newaxis], -uvecs, uvecs)

    # Normalize eigenvectors to point "up" (positive in final axis)
    down_mask = uvecs[..., -1] < 0  # (..., n)
    uvecs = np.where(down_mask[..., np.newaxis], -uvecs, uvecs)

    # Sort by magnitude (largest first)
    idxs = np.argsort(-np.abs(vals), axis=-1)  # (..., n)
    vals = np.take_along_axis(vals, idxs, axis=-1)
    uvecs = np.take_along_axis(uvecs, idxs[..., np.newaxis], axis=-2)

    # Return single matrix result without batch dimension
    if single:
        return Decomposition(vals[0], uvecs[0])
    return Decomposition(vals, uvecs)


def is_symmetric(
    mat: NDArray[np.floating],
) -> bool | NDArray[np.bool_]:
    """Check if a matrix (or batch of matrices) is exactly symmetric.

    Args:
        mat: A single matrix of shape (n, n) or batched matrices of shape (..., n, n).

    Returns:
        For single matrix: bool indicating if symmetric.
        For batched matrices: boolean array of shape (...) indicating symmetry of each.
        Returns False if input is not a valid square matrix shape.
    """
    if mat.ndim < 2 or mat.shape[-1] != mat.shape[-2]:  # noqa: PLR2004
        return False  # Not a square matrix
    transposed = np.swapaxes(mat, -1, -2)
    result = np.all(mat == transposed, axis=(-1, -2))
    # Return scalar bool for single matrix, array for batched
    return bool(result) if result.ndim == 0 else result
