from itertools import groupby
from typing import Literal

import numpy as np
from numpy.typing import NDArray

from qoptcraft.algebra import (
    hermitian_matrix_from_coefs,
    hilbert_dim,
    image_algebra_basis,
    unitary_algebra_basis,
)
from qoptcraft.invariant import (
    operator_higher_order_projection,
    operator_nested_commutator,
)
from qoptcraft.math import hs_inner_product

# from qoptcraft.testing import assert_hermitian, assert_symmetric
from qoptcraft.utils import saved_basis


@saved_basis(file_name="recursive_subspaces")
def recursive_invariant_subspaces(
    modes,
    photons,
    *,
    invariant_operator: Literal["higher_order_projection", "nested_commutator"],
    order: int,
    orthonormal: bool = False,
    cache: bool = True,
    parallelize=True,
):

    _ = cache  # only used by the decorator @saved_basis

    if invariant_operator not in [
        "higher_order_projection",
        "nested_commutator",
    ]:
        raise ValueError(f"Invalid {invariant_operator = }")

    is_nested_commutator = invariant_operator == "nested_commutator"

    if (is_nested_commutator and order <= 2) or order <= 1:
        subspace_partition = [None]
    else:
        subspace_partition = recursive_invariant_subspaces(
            modes,
            photons,
            invariant_operator=invariant_operator,
            order=order - 1,
            orthonormal=orthonormal,
            cache=cache,
            parallelize=parallelize,
        )

    finer_subspace_partition = []
    for subspace in subspace_partition:
        finer_subspace_partition += invariant_subspaces(
            modes,
            photons,
            invariant_operator=invariant_operator,
            order=order,
            orthonormal=orthonormal,
            cache=False,
            basis=subspace,
            parallelize=parallelize,
        )
    return finer_subspace_partition


@saved_basis(file_name="subspaces")
def invariant_subspaces(
    modes,
    photons,
    *,
    invariant_operator: Literal["higher_order_projection", "nested_commutator"],
    order: int,
    orthonormal: bool = False,
    cache: bool = True,
    parallelize=True,
    basis: list[NDArray] | None = None,
) -> list[list[NDArray]]:
    _ = cache  # only used by the decorator @saved_basis

    if invariant_operator not in [
        "higher_order_projection",
        "nested_commutator",
    ]:
        raise ValueError(f"Invalid {invariant_operator = }")

    operator = self_adjoint_projection(
        modes,
        photons,
        invariant_operator=invariant_operator,
        order=order,
        orthonormal=orthonormal,
        basis=basis,
        parallelize=parallelize,
    )

    eigenvalues, eigenvectors = np.linalg.eigh(operator)
    eigenvectors[:, :] = eigenvectors.T  # eigenvectors will be rows instead of columns
    eigenvalues[:], eigenvectors[:] = zip(
        *sorted(zip(eigenvalues.round(4), eigenvectors, strict=True), key=lambda p: p[0]),
        strict=True,
    )

    # print(f"\nThe eigenvalues for {basis = } are {eigenvalues}\n")

    # sorted_pairs = sorted(zip(eigenvalues, eigenvectors), key=lambda p: round(p[0], 6))
    # eigenvalues = np.array([round(ev, 6) for ev, _ in sorted_pairs])
    # eigenvectors = np.array([vec for _, vec in sorted_pairs])

    return [
        [hermitian_matrix_from_coefs(c, basis) for _, c in g]
        for _, g in groupby(zip(eigenvalues, eigenvectors, strict=True), key=lambda x: x[0])
    ]


def self_adjoint_projection(
    modes,
    photons,
    invariant_operator: Literal["higher_order_projection", "nested_commutator"],
    order: int,
    orthonormal=False,
    parallelize=True,
    basis: list[NDArray] | None = None,
) -> NDArray:
    dim = hilbert_dim(modes, photons)

    if invariant_operator not in [
        "higher_order_projection",
        "nested_commutator",
    ]:
        raise ValueError(f"Invalid {invariant_operator = }")

    is_nested_commutator = invariant_operator == "nested_commutator"

    if is_nested_commutator:
        scattering_basis = unitary_algebra_basis(modes)
    image_basis = image_algebra_basis(modes, photons, orthonormal=orthonormal)

    if basis is None:
        basis = unitary_algebra_basis(dim)

    operator_matrix = np.zeros((len(basis), len(basis)), dtype=np.complex128)

    if parallelize:
        from functools import partial

        from pathos.multiprocessing import ProcessingPool, cpu_count

        compute_row = partial(
            row_self_adjoint_projection,
            basis=basis,
            invariant_operator=invariant_operator,
            modes=modes,
            order=order,
            image_basis=image_basis,
            scattering_basis=scattering_basis if is_nested_commutator else None,
        )

        with ProcessingPool(nodes=cpu_count()) as pool:
            parallelized_rows = pool.map(compute_row, range(len(basis)))
        for i, row in parallelized_rows:
            operator_matrix[i, :] = row
    else:
        for i in range(len(basis)):
            operator_matrix[i, :] = row_self_adjoint_projection(
                i,
                basis=basis,
                invariant_operator=invariant_operator,
                modes=modes,
                order=order,
                image_basis=image_basis,
                scattering_basis=scattering_basis if is_nested_commutator else None,
            )[
                1
            ]  # first output is just i

    for i in range(len(basis)):
        operator_matrix[i, :] = row_self_adjoint_projection(
            i,
            basis=basis,
            invariant_operator=invariant_operator,
            modes=modes,
            order=order,
            image_basis=image_basis,
            scattering_basis=scattering_basis if is_nested_commutator else None,
        )[
            1
        ]  # first output is just i

    return operator_matrix  # small errors in the imaginary part give problems


def row_self_adjoint_projection(
    i,
    basis,
    invariant_operator: Literal["higher_order_projection", "nested_commutator"],
    modes,
    order,
    image_basis,
    scattering_basis,
):
    if invariant_operator == "higher_order_projection":
        projected = operator_higher_order_projection(basis[i], modes, order, image_basis)
    elif invariant_operator == "nested_commutator":
        projected = operator_nested_commutator(basis[i], order, scattering_basis, image_basis)
    else:
        raise ValueError(f"Invalid invariant_operator: {invariant_operator}")

    # ! if hs_inner_product not casted into real, numerical problems will arise
    # ! the coefficients are real because the vector space is real
    row = np.array([hs_inner_product(matrix, projected).real for matrix in basis])
    return i, row
