"""Helpers for exact Schmidt spectra and entanglement calculations."""

import numpy as np
from scipy.linalg import eigh as array_eigh

__all__ = ["dense_schmidt_probabilities"]


def dense_schmidt_probabilities(psi_matrix: np.ndarray) -> np.ndarray:
    """Return the exact Schmidt probabilities of a dense bipartition matrix."""
    n_rows, n_cols = psi_matrix.shape
    if n_rows <= n_cols:
        gram_matrix = psi_matrix @ psi_matrix.conj().T
    else:
        gram_matrix = psi_matrix.conj().T @ psi_matrix
    llambdas = array_eigh(gram_matrix, eigvals_only=True)
    return np.maximum(llambdas, 0.0)
