import numpy as np
from numpy.typing import NDArray

from qoptcraft.algebra import image_algebra_basis
from qoptcraft.state import PureState


def covariance_matrix(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M_ij = Tr(O_i rho)Tr(O_j rho) - 0.5 Tr(O_iO_j + O_jO_i rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)
    dim = len(image_basis)
    matrix = np.zeros((dim, dim), dtype=np.complex128)
    for i, basis_i in enumerate(image_basis):
        for j, basis_j in enumerate(image_basis):
            matrix[i, j] = np.trace(basis_i @ state.density_matrix) * np.trace(
                basis_j @ state.density_matrix
            )
            matrix[i, j] -= 0.5 * np.trace(
                (basis_i @ basis_j + basis_j @ basis_i) @ state.density_matrix
            )
    return matrix


def covariance_invariant(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M_ij = Tr(O_i rho)Tr(O_j rho) - 0.5 Tr(O_iO_j + O_jO_i rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    return np.linalg.eigvalsh(covariance_matrix(state, orthonormal)).round(23)


def m1_matrix(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M1 = Tr(O_i rho)Tr(O_j rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    dim = len(image_basis)
    invariant = np.zeros((dim, dim), dtype=np.complex128)
    for i, basis_i in enumerate(image_basis):
        for j, basis_j in enumerate(image_basis):
            invariant[i, j] = np.trace(basis_i @ state.density_matrix) * np.trace(
                basis_j @ state.density_matrix
            )

    return invariant


def m1_invariant(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M1 = Tr(O_i rho)Tr(O_j rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    dim = len(image_basis)
    invariant = np.zeros((dim, dim), dtype=np.complex128)
    for i, basis_i in enumerate(image_basis):
        for j, basis_j in enumerate(image_basis):
            invariant[i, j] = np.trace(basis_i @ state.density_matrix) * np.trace(
                basis_j @ state.density_matrix
            )

    return np.linalg.eigvals(invariant).round(23)


def m2_matrix(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M2 = Tr(O_iO_j + O_jO_i rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    dim = len(image_basis)
    invariant = np.zeros((dim, dim), dtype=np.complex128)
    for i, basis_i in enumerate(image_basis):
        for j, basis_j in enumerate(image_basis):
            invariant[i, j] = 0.5 * np.trace(
                (basis_i @ basis_j + basis_j @ basis_i) @ state.density_matrix
            )

    return invariant


def m2_invariant(state: PureState, orthonormal=False) -> NDArray:
    """Calculate M2 = Tr(O_iO_j + O_jO_i rho).

    Args:
        state (State): a photonic quantum state.

    Returns:
        NDArray: spectrum of the total matrix.
    """
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    dim = len(image_basis)
    invariant = np.zeros((dim, dim), dtype=np.complex128)
    for i, basis_i in enumerate(image_basis):
        for j, basis_j in enumerate(image_basis):
            invariant[i, j] = 0.5 * np.trace(
                (basis_i @ basis_j + basis_j @ basis_i) @ state.density_matrix
            )

    return np.linalg.eigvals(invariant).round(23)
