"""Density-matrix helpers used by entanglement and LBO workflows."""

import logging
import numpy as np
from scipy.linalg import eigh as array_eigh
from scipy.sparse import csr_matrix, isspmatrix

logger = logging.getLogger(__name__)

__all__ = [
    "diagonalize_density_matrix",
    "get_projector_for_efficient_density_matrix",
]


def diagonalize_density_matrix(rho: np.ndarray | csr_matrix):
    """Diagonalize a dense or sparse Hermitian density matrix."""
    if isinstance(rho, np.ndarray):
        rho_eigvals, rho_eigvecs = array_eigh(rho)
    elif isspmatrix(rho):
        rho_eigvals, rho_eigvecs = array_eigh(rho.toarray())
    else:
        raise TypeError("rho must be a NumPy array or SciPy sparse matrix.")
    return rho_eigvals, rho_eigvecs


def get_projector_for_efficient_density_matrix(
    rho_eigvals, rho_eigvecs, threshold: float
):
    """Build a projector from the dominant eigenvectors of a density matrix."""
    sorted_indices = np.argsort(rho_eigvals)[::-1]
    rho_eigvals = rho_eigvals[sorted_indices]
    rho_eigvecs = rho_eigvecs[:, sorted_indices]
    p_columns = np.sum(rho_eigvals > threshold)
    while p_columns < 2:
        threshold /= 10
        p_columns = np.sum(rho_eigvals > threshold)
    logger.info(f"SIGNIFICANT EIGENVALUES {p_columns} with threshold {threshold}")
    proj = np.zeros((rho_eigvals.shape[0], p_columns), dtype=complex)
    for jj in range(p_columns):
        proj[:, jj] = rho_eigvecs[:, jj]
    return proj
