from itertools import product
from typing import Literal

import numpy as np
from numpy.typing import NDArray

from qoptcraft.algebra import image_algebra_basis, unitary_algebra_basis, projection_subspace
from qoptcraft.math import Matrix, gram_schmidt, hs_scalar_product
from qoptcraft.state import State


def projection_density(
    state: State,
    subspace: list[NDArray] | Literal["preimage", "image", "complement", "full"] = "preimage",
    orthonormal: bool = False,
) -> Matrix:
    """Project a state onto the linear optical subalgebra, its complement or the preimage.

    Args:
        state (State): a photonic quantum state.
        subspace (str): "image" or "complement".
        orthonormal (bool, optional): if true, it orthonormalizes the basis of linear optical
            hamiltonians before projecting onto it. Defaults to False.

    Returns:
        Matrix: the projected density matrix.
    """
    matrix = state.density_matrix

    if isinstance(subspace, list):
        if orthonormal:
            subspace = gram_schmidt(subspace)
        return projection_subspace(matrix, subspace)

    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)
    projection_image = projection_subspace(matrix, image_basis)
    if subspace == "image":
        return projection_image
    if subspace == "complement" and orthonormal is True:
        return matrix - projection_image
    if subspace == "preimage":
        unitary_basis = unitary_algebra_basis(state.modes)
        projection = np.zeros_like(unitary_basis[0])
        for unitary_matrix, image_matrix in zip(unitary_basis, image_basis, strict=True):
            projection += hs_scalar_product(image_matrix, matrix) * unitary_matrix
        return projection
    raise ValueError("The complement invariant only makes sense if the basis is orthonormal.")


def higher_order_projection_density(
    state: State,
    order: int,
    subspace: Literal["preimage", "image"] = "image",
    orthonormal: bool = False,
) -> Matrix:
    """Project a state onto Σ tr(O_i ... O_k ρ) O_i ... O_k with subspace="image"
    or Σ tr(O_i ... O_k ρ) b_i ... b_k with subspace="preimage"

    Args:
        state (State): a photonic quantum state.
        subspace (str): "image" or "complement".
        orthonormal (bool, optional): if true, it orthonormalizes the basis of linear optical
            hamiltonians before projecting onto it. Defaults to False.

    Returns:
        Matrix: the projected density matrix.
    """
    density_matrix = state.density_matrix

    unitary_basis = unitary_algebra_basis(state.modes)
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    if subspace == "preimage":
        if orthonormal:
            raise ValueError("Orthonormal basis doesnt work with the preimage subspace.")
        projection = np.zeros_like(unitary_basis[0])
        for indices in product(*[list(range(state.modes * state.modes))] * order):
            hermitian_matrix = unitary_basis[indices[0]]
            image_matrix = image_basis[indices[0]]
            for idx in indices[1:]:
                hermitian_matrix = hermitian_matrix @ unitary_basis[idx]
                image_matrix = image_matrix @ image_basis[idx]
            projection = (
                projection + hs_scalar_product(image_matrix, density_matrix) * hermitian_matrix
            )
    elif subspace == "image":
        projection = np.zeros(image_basis[0].shape, dtype=np.complex128)
        for indices in product(*[list(range(state.modes * state.modes))] * order):
            image_matrix = image_basis[indices[0]]
            for idx in indices[1:]:
                image_matrix = image_matrix @ image_basis[idx]
            projection = projection + hs_scalar_product(image_matrix, density_matrix) * image_matrix
    else:
        raise ValueError(f"Only admited subspaces are preimage and image, but {subspace = }.")
    return projection
