import numpy as np
from numpy.typing import NDArray

from qoptcraft.algebra import image_algebra_basis
from .basis import unitary_algebra_basis
from qoptcraft.math import hs_scalar_product
from .utils import get_coordinates


# TODO: consider sparse matrices
def algebra_hm(matrix: NDArray, photons: int, orthonormal: bool = False, sparse = True) -> NDArray:
    """Takes a hermitian matrix and outputs its image under the differential
    of the photonic homomorphism.

    Args:
        matrix (_type_): _description_
        orthonormal (bool, optional): Orthonormal image basis. Defaults to False.

    Returns:
        _type_: _description_
    """
    modes = matrix.shape[1]

    basis = unitary_algebra_basis(modes)
    image_basis = image_algebra_basis(modes, photons, orthonormal=orthonormal, sparse=sparse)

    image_matrix = hs_scalar_product(matrix, basis[0]) * image_basis[0]

    for bi, Oi in zip(basis[1:], image_basis[1:], strict=True):
        image_matrix = image_matrix + hs_scalar_product(matrix, bi) * Oi
    return image_matrix


def inv_algebra_hm(matrix, modes, photons, change_basis=None):
    """Project a matrix into the image subalgebra and compute the inverse algebra homomorphism."""

    if change_basis is None:
        change_basis = image_orthonormal_to_standard_basis(modes, photons)

    projection_image = np.zeros((modes, modes), dtype=np.complex128)
    orthonormal_basis = image_algebra_basis(modes, photons, orthonormal=True)
    image_coefs = get_coordinates(matrix, orthonormal_basis)

    for i, bi in enumerate(unitary_algebra_basis(modes)):
        projection_image += np.dot(image_coefs, change_basis[:, i]) * bi
    return projection_image


def image_orthonormal_to_standard_basis(modes, photons):

    image_basis = image_algebra_basis(modes, photons, orthonormal=False)
    orthonormal_basis = image_algebra_basis(modes, photons, orthonormal=True)

    C = np.zeros((modes**2, modes**2))
    for i in range(modes**2):
        for j in range(modes**2):
            C[i, j] = hs_scalar_product(image_basis[i], orthonormal_basis[j])
    return np.linalg.inv(C)
