"""Vector normalisation and symmetric orthogonalisation."""
from __future__ import annotations

import numpy as np

EPS = 1e-9


def unit(a: np.ndarray) -> np.ndarray:
    """Rows of `a` scaled to unit L2 norm, in float32."""
    a = np.asarray(a, np.float32)
    return a / (np.linalg.norm(a, axis=-1, keepdims=True) + EPS)


def lowdin(M: np.ndarray) -> np.ndarray:
    """Symmetric (Loewdin) orthonormalisation U = M (M^T M)^{-1/2}.

    Each output column is sign-aligned with the prototype column it came from, so axis k of
    the result still points the same way as column k of `M`.
    """
    M = np.asarray(M, np.float32)
    w, V = np.linalg.eigh(M.T @ M)
    U = M @ (V @ np.diag(np.maximum(w, 1e-12) ** -0.5) @ V.T)
    return U * np.sign(np.einsum("hc,hc->c", U, M))[None, :]