"""Stable Gram-Schmidt algorithm."""

from collections.abc import Generator

from numpy.typing import NDArray
from scipy.sparse import sparray

from .norms import hs_inner_product, hs_norm


def gram_schmidt(basis: list[sparray] | list[NDArray]) -> list[sparray] | list[NDArray]:
    """Gram-Schmidt algorithm to orthonormalize a basis.

    Note:
        It turns out that the normal Gram-Schmidt algorithm suffers from
        numerical instability: Round-off errors can accumulate and destroy
        orthogonality of the resulting vectors. We introduce the modified
        Gram-Schmidt procedure to help remedy this issue.

    Args:
        basis (list[sparray]): basis to orthonormalize.

    Returns:
        list[sparray]: orthonormalized basis.

    References:
        Algorithm can be found in https://www.math.uci.edu/~ttrogdon/105A/html/Lecture23.html
    """
    dim = len(basis)
    orth_basis = []

    for j in range(dim):
        orth_basis.append(basis[j] / hs_norm(basis[j]))
        for k in range(j + 1, dim):
            basis[k] = basis[k] - hs_inner_product(orth_basis[j], basis[k]) * orth_basis[j]
    return orth_basis


def gram_schmidt_generator(basis: list[sparray] | list[NDArray]) -> Generator:
    """Gram-Schmidt algorithm to orthonormalize a basis.

    Note:
        It turns out that the normal Gram-Schmidt algorithm suffers from
        numerical instability: Round-off errors can accumulate and destroy
        orthogonality of the resulting vectors. We introduce the modified
        Gram-Schmidt procedure to help remedy this issue.

    Args:
        basis (list[sparray]): basis to orthonormalize.

    Yields:
        NDArray or sparray: orthonormalized basis element.

    References:
        Algorithm can be found in https://www.math.uci.edu/~ttrogdon/105A/html/Lecture23.html
    """
    dim = len(basis)
    orth_basis = []

    for j in range(dim):
        orth_basis.append(basis[j] / hs_norm(basis[j]))
        yield orth_basis[j]
        for k in range(j + 1, dim):
            basis[k] = basis[k] - hs_inner_product(orth_basis[j], basis[k]) * orth_basis[j]
