from itertools import product

import numpy as np

from qoptcraft.algebra import hilbert_dim, image_algebra_basis, unitary_algebra_basis
from qoptcraft.math import Matrix, commutator, hs_inner_product
from qoptcraft.state import State
from qoptcraft.testing import assert_hermitian


def invariant_coef(indices, basis):
    """Compute tr(b_i ... b_k)"""
    matrix = np.eye(basis[0].shape[0], dtype=np.complex128)
    for i in indices:
        matrix @= basis[i]
    return np.trace(matrix)


def equivalent_spectral_invariant(state: State, order: int, orthonormal=False, image=False) -> Matrix:
    """Theorem 4 of "Lie algebraic invariants" invariants. They are equivalent to the
    tangent projection invariant."""
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)
    coef_basis = image_basis if image else unitary_algebra_basis(state.modes)

    invariant = 0

    for indices in product(*[list(range(state.modes**2))] * order):
        coef = invariant_coef(indices, coef_basis)
        if coef != 0:
            for idx in indices:
                coef *= np.trace(image_basis[idx] @ state.density_matrix)
            invariant += coef
    return np.real(invariant)


def casimir_operator(modes: int, photons: int, order: int, orthonormal=False) -> Matrix:
    """These operators are multiples of the identity."""

    scattering_basis = unitary_algebra_basis(modes)
    image_basis = image_algebra_basis(modes, photons, orthonormal=orthonormal)

    dim = hilbert_dim(modes, photons)
    invariant = np.zeros((dim, dim), dtype=np.complex128)

    for indices in product(*[list(range(modes * modes))] * order):
        matrix = np.eye(dim, dtype=np.complex128)
        coef = invariant_coef(indices, scattering_basis)
        if coef != 0:
            for idx in indices:
                matrix @= image_basis[idx]
            invariant += coef * matrix

    return invariant


def invariant_operator_commutator(state: State, order: int, orthonormal=False) -> Matrix:
    """Compute the spectrum of Σ tr(b_i ... b_k) [O_i, ρ] ... [O_k, ρ]"""
    scattering_basis = unitary_algebra_basis(state.modes)
    image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

    dim = hilbert_dim(state.modes, state.photons)
    invariant = np.zeros((dim, dim), dtype=np.complex128)

    for indices in product(*[list(range(state.modes * state.modes))] * order):
        matrix = np.eye(dim, dtype=np.complex128)
        coef = invariant_coef(indices, scattering_basis)
        if coef != 0:
            for idx in indices:
                matrix @= commutator(image_basis[idx], state.density_matrix)
            invariant += coef * matrix

    return np.linalg.eigvals(invariant)


# def invariant_operator_traces(state: State, order: int, orthonormal=False) -> Matrix:
#     scattering_basis = unitary_algebra_basis(state.modes)
#     image_basis = image_algebra_basis(state.modes, state.photons, orthonormal=orthonormal)

#     dim = hilbert_dim(state.modes, state.photons)
#     invariant = np.zeros((dim, dim), dtype=np.complex128)

#     for indices in product(*[list(range(state.modes * state.modes))] * order):
#         matrix = np.eye(dim, dtype=np.complex128)
#         coef = invariant_coef(indices, scattering_basis)
#         if coef != 0:
#             for idx in indices:
#                 matrix @= image_basis[idx] @ state.density_matrix
#             invariant += coef * matrix

#     traces = []
#     for i in range(1, dim + 1):
#         traces.append(np.trace(np.linalg.matrix_power(invariant, i)))

#     return traces


def operator_higher_order_projection(matrix, modes: int, order: int, image_basis):
    """Theorem 7 of "Lie algebraic invariants" higher order projection operator."""
    projection = np.zeros(image_basis[0].shape, dtype=np.complex128)
    for indices in product(*[list(range(modes * modes))] * order):
        image_matrix = image_basis[indices[0]]
        for idx in indices[1:]:
            image_matrix = image_matrix @ image_basis[idx]
        projection += hs_inner_product(image_matrix, matrix) * image_matrix
    assert_hermitian(projection)
    return projection


def operator_nested_commutator(matrix, order: int, scattering_basis, image_basis):
    """Theorem 8 of "Lie algebraic invariants" nested commutator operator N_k."""

    modes = scattering_basis[0].shape[0]

    mapped_matrix = np.zeros(image_basis[0].shape, dtype=np.complex128)

    # for indices in combinations_with_replacement(range(modes * modes), order):
    #     coef = symmetrized_invariant_coef(indices, scattering_basis)
    #     if coef != 0:
    #         for perm_indices in permutations(indices):
    #             temp_matrix = commutator(image_basis[perm_indices[0]], matrix)
    #             for idx in perm_indices[1:]:
    #                 temp_matrix = commutator(image_basis[idx], temp_matrix)
    #             mapped_matrix += coef * temp_matrix

    for indices in product(*[list(range(modes * modes))] * order):
        coef = invariant_coef(indices, scattering_basis)
        if coef != 0:
            temp_matrix = commutator(image_basis[indices[0]], matrix)
            for idx in indices[1:]:
                temp_matrix = commutator(image_basis[idx], temp_matrix)
            mapped_matrix += coef * temp_matrix

    mapped_matrix *= (-1) ** order
    assert_hermitian(mapped_matrix)

    return mapped_matrix


def invariant_nested_commutator(state: State, order: int, orthonormal=False):
    """Theorem 8 of "Lie algebraic invariants" nested commutator invariant."""

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

    invariant = operator_nested_commutator(
        state.density_matrix, order, scattering_basis, image_basis
    )

    return np.sort(np.linalg.eigvals(invariant))
