from typing import Literal
# from permanent import permanent_ryser as permanent  # efficient rust package for permanents

import numpy as np
from numpy.typing import NDArray
from scipy.special import factorial as factorial_array
import scipy as sp
import numba

from qoptcraft.state import Fock, PureState
from qoptcraft.algebra import photon_basis, BasisPhoton, algebra_hm
from qoptcraft.math import permanent, log_matrix


def factorial_cache(n: int) -> NDArray:
    fact = np.empty(n + 1)
    fact[0] = 1
    for i in range(1, n + 1):
        fact[i] = fact[i - 1] * i
    return fact


@numba.jit(nopython=True)
def fact_coef(fock_out, fact: NDArray) -> int:
    coef_out = 1.0
    for n in fock_out:
        coef_out *= fact[n]
    return coef_out


def fock_evolution(
    scattering_matrix: NDArray,
    fock_in: tuple[int, ...],
    method: Literal["hamiltonian", "heisenberg", "permanent glynn", "permanent ryser"] = "permanent glynn",
) -> NDArray:
    """Evolution of a single Fock state using the definition given by basic
    quantum mechanics.

    Args:
        scattering_matrix (NDArray): scattering matrix of a linear optical interferometer.
        fock_in (Fock): fock state to evolve.
        method (str): method to calculate the evolution of the Fock state. Options are
            'heisenberg', 'permanent glynn' or 'permanent ryser'. Default is 'permanent glynn'.

    Returns:
        NDArray: fock state given in the photon basis.
    """
    if method == "hamiltonian":
        return fock_evolution_hamiltonian(scattering_matrix, fock_in)
    if method == "heisenberg":
        return fock_evolution_heisenberg(scattering_matrix, fock_in)
    if method.split()[0] == "permanent":
        return fock_evolution_permanent(scattering_matrix, fock_in, method=method.split()[1])
    raise ValueError("Options for method are 'heisenberg', 'permanent glynn' or 'permanent ryser'.")


def fock_evolution_hamiltonian(scattering_matrix: NDArray, fock_in: tuple[int, ...]) -> NDArray:
    """ "Unitary matrix of a linear interferometer with a number of photons as input.
    Calculated by evolving the hamiltonian of the interferometer.

    Args:
        scattering_matrix (NDArray): scattering matrix of a linear optical interferometer.
        photons (int): number of photons.

    Returns:
        NDArray: image of the scattering matrix through the photonic homomorphism.
    """
    photons = sum(fock_in)
    state_vec = Fock(*fock_in).state_vector
    S_hamiltonian = -1j * log_matrix(scattering_matrix, method="schur")  # ? Best method??
    U_hamiltonian = algebra_hm(S_hamiltonian, photons)
    return sp.sparse.linalg.expm_multiply(1j * U_hamiltonian, state_vec)


def fock_evolution_heisenberg(scattering_matrix: NDArray, fock_in: tuple[int, ...]) -> NDArray:
    """Evolution of a single Fock state using the definition given by basic
    quantum mechanics.

    Args:
        scattering_matrix (NDArray): scattering matrix of a linear optical interferometer.
        fock_in (Fock): fock state to evolve.

    Returns:
        NDArray: fock state given in the photon basis.
    """
    modes = len(fock_in)

    state_out: PureState = Fock(*([0] * modes))
    for mode_2 in range(modes):
        for _ in range(fock_in[mode_2]):
            state_aux = state_out
            state_out = scattering_matrix[0, mode_2] * state_aux.creation(0)
            for mode_1 in range(1, modes):
                state_out += scattering_matrix[mode_1, mode_2] * state_aux.creation(mode_1)
    coef = np.prod(factorial_array(fock_in))
    return state_out / np.sqrt(coef)


def fock_evolution_permanent(
    scattering_matrix: NDArray,
    fock_in: tuple[int, ...],
    method: Literal["glynn", "ryser"] = "glynn",
    photonic_basis: BasisPhoton = None,
) -> NDArray:
    """Evolution of a single Fock state using the definition given by basic
    quantum mechanics.

    Args:
        scattering_matrix (NDArray): scattering matrix of a linear optical interferometer.
        fock_in (Fock): fock state to evolve.
        method (str): method to compute the permanent. Must be 'glynn' or 'ryser'.
            Defaults to 'glynn'.

    Returns:
        NDArray: fock state given in the photon basis.
    """
    if len(fock_in) != scattering_matrix.shape[0]:
        raise ValueError("Dimension of scattering_matrix and number of modes don't match.")
    if photonic_basis is None:
        photonic_basis = photon_basis(len(fock_in), sum(fock_in))

    modes = len(fock_in)
    photons = sum(fock_in)
    dim = len(photonic_basis)

    state_out = np.empty(dim, dtype=np.complex128)
    sub_matrix = np.empty((photons, photons), dtype=np.complex128)

    fact = factorial_cache(photons)
    for row, fock_out in enumerate(photonic_basis):
        in_out_submatrix_inplace(
            sub_matrix, scattering_matrix, fock_in, fock_out, modes
        )
        coef = fact_coef(fock_out, fact)
        state_out[row] = permanent(sub_matrix, method=method) / np.sqrt(coef)
        # state_out[row] = permanent(sub_matrix) / np.sqrt(coef)
    coef = fact_coef(fock_in, fact)
    return state_out / np.sqrt(coef)


# def _fock_evolution_permanent(
#     scattering_matrix: NDArray,
#     fock_in: tuple[int, ...],
#     method: Literal["glynn", "ryser"],
#     photonic_basis: BasisPhoton,
#     fact_cache: NDArray
# ) -> NDArray:
#     ...


@numba.jit(nopython=True)
def in_out_submatrix_inplace(
    sub_matrix,
    scattering_matrix,
    fock_in,
    fock_out,
    modes,
):
    """
    Fill sub_matrix in-place using input/output Fock occupations.

    Args:
        sub_matrix (n, n): preallocated output matrix
        matrix (m, m): scattering matrix
        fock_in (m,): input occupation numbers
        fock_out (m,): output occupation numbers
        modes (int): number of modes
    """
    col = 0
    for in_mode in range(modes):
        for _ in range(fock_in[in_mode]):
            row = 0
            for out_mode in range(modes):
                for _ in range(fock_out[out_mode]):
                    sub_matrix[row, col] = scattering_matrix[out_mode, in_mode]
                    row += 1
            col += 1


# @numba.jit(nopython=True)
# def in_out_submatrix_inplace(
#     sub_matrix, interm_matrix, matrix, fock_in, fock_out, modes, photons
# ) -> NDArray:
#     """Return a matrix with row index 'i' repeated 'row_rep' times
#     and column index 'j' repeated 'col_rep' times.

#     Args:
#         matrix ((m, m) array): Linear optical scattering matrix with m modes.

#     Returns:
#         (n, n) array: Optical scattering matrix with m modes and n photons.
#     """
#     col = 0
#     for mode in range(modes):
#         for _ in range(fock_in[mode]):
#             for i in range(modes):
#                 interm_matrix[i, col] = matrix[i, mode]
#             col += 1
#     row = 0
#     for mode in range(modes):
#         for _ in range(fock_out[mode]):
#             for j in range(photons):
#                 sub_matrix[row, j] = interm_matrix[mode, j]
#             row += 1
