import itertools
from collections.abc import Callable, Mapping, Sequence
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import numpy.typing as npt
from deltakit_circuit._circuit import Circuit
from deltakit_decode._mwpm_decoder import PyMatchingDecoder
from deltakit_decode.analysis._matching_decoder_managers import StimDecoderManager
from tqdm import tqdm

from deltakit_explorer.analysis.error_budget._memory import (
    MemoryGenerator,
    get_rotated_surface_code_memory_circuit,
)


def _generate_surface_code_memory_decoder_manager(
    distance: int,
    num_rounds: int,
    noise_model: Callable[[Circuit, npt.NDArray[np.floating]], Circuit],
    noise_parameters: npt.NDArray[np.floating],
    memory_generator: MemoryGenerator,
    noise_parameter_names: Sequence[str],
) -> StimDecoderManager:
    """Generate a decoder manager with a rotated planar code memory experiment.

    Args:
        distance (int): distance of the rotated planar code memory experiment that
            should be generated by this function.
        num_rounds (int): number of stabilizer measurement rounds of the memory
            experiment generated by this function.
        noise_model (Callable[[Circuit, npt.NDArray[np.floating]], Circuit]): a callable
            adding noise to the provided circuit, according to the parameters provided
            in the second argument.
        noise_parameters: parameters for the ``noise_model`` callable representing the
            noise strength.
        memory_generator: a callable that is able to generate memory experiments from a
            given distance and number of rounds.
        noise_parameter_names: if provided, human-readable names for each of the
            provided ``noise_parameters``.

    Returns:
        a decoder manager that can then be run using ``RunAllAnalysisEngine``.
    """
    circuit = memory_generator(distance, num_rounds)
    noisy_circuit = noise_model(circuit, noise_parameters)
    decoder, decoder_circuit = PyMatchingDecoder.construct_decoder_and_stim_circuit(
        noisy_circuit
    )
    metadata = {
        "distance": distance,
        "num_rounds": num_rounds,
        **{
            f"noise_{name}": p
            for name, p in zip(noise_parameter_names, noise_parameters, strict=True)
        },
    }

    return StimDecoderManager(decoder_circuit, decoder, metadata=metadata)


def _generate_surface_code_memory_decoder_manager_wrapper(
    data: tuple[
        int,
        int,
        Callable[[Circuit, npt.NDArray[np.floating]], Circuit],
        npt.NDArray[np.floating],
        MemoryGenerator,
        Sequence[str],
    ],
) -> StimDecoderManager:
    return _generate_surface_code_memory_decoder_manager(*data)


def generate_decoder_managers_for_lambda(
    xi: npt.NDArray[np.floating],
    noise_model: Callable[[Circuit, npt.NDArray[np.floating]], Circuit],
    num_rounds_by_distances: Mapping[int, Sequence[int]],
    max_workers: int = 1,
    memory_generator: MemoryGenerator = get_rotated_surface_code_memory_circuit,
    noise_parameter_names: Sequence[str] | None = None,
) -> list[StimDecoderManager]:
    """Generate several decoder managers from the provided arguments for a rotated
    planar code memory experiment.

    This function should be used to generate in one call a description of all the
    sampling tasks that should be done in order to estimate the value of Λ on the
    provided noise parameters ``xi``.

    Note:
        Generating the tasks might be a long process, hence the presence of the
        ``max_worker`` parameter.

    Warning:
        This function will generate ``len(num_rounds) * len(distances) *
        numpy.prod(xi.shape[1:])`` sampling tasks. That quantity
        might grow quickly, so try to avoid providing too many entries in ``distances``
        and ``num_rounds`` except if you really need it.

    Args:
        xi (npt.NDArray[numpy.floating]): an array of shape ``(m, n)``, where ``m`` is
            the number of noise parameters needed by ``noise_model`` and ``n`` the
            number of points to evaluate Λ on.
        noise_model (Callable[[Circuit, npt.NDArray[np.floating]], Circuit]): a callable
            adding noise to the provided circuit, according to the parameters provided.
        num_rounds_by_distances (Mapping[int, Sequence[int]]): the different number of
            rounds that will be used to estimate the logical error-rate per round for
            each code distance. As a rule of thumb, at least ``4`` entries should be
            provided and the maximum number of rounds should, if possible, give a
            logical error-rate above ``0.2`` for the largest distance and a logical
            error-rate below ``0.45`` for the smallest distance. You will not be able to
            check those thresholds everytime, but the standard errors computed will be
            lower if they are verified.
        max_workers (int): if provided, the number of parallel processes that are used
            to generate the sampling tasks. If not, default to ``1``.
        memory_generator: a callable that is able to generate memory experiments from a
            given distance and number of rounds.
        noise_parameter_names: if provided, human-readable names for each of the
            provided ``noise_parameters``. Defaults to the noise parameter index (i.e.,
            "0", "1", ...).

    Returns:
        a list of decoder managers that can then be run using ``RunAllAnalysisEngine``.
    """
    if noise_parameter_names is None:
        noise_parameter_names = [str(i) for i in range(xi.shape[0])]
    # 0. Transposing to be able to iterate on different values.
    noise_parameters = xi.T
    # 1. Check input parameters
    if max_workers <= 0:
        msg = f"Cannot have less than one worker. Asked for {max_workers} workers."
        raise ValueError(msg)
    # 2. Build an iterator on all the parameters we need to provide to the generation
    # method.
    total_circuits = (
        sum(len(num_rounds) for num_rounds in num_rounds_by_distances.values())
        * noise_parameters.shape[0]
    )
    # (distance: int, num_rounds: int)
    distance_and_rounds_iterator = itertools.chain.from_iterable(
        zip(itertools.repeat(d, len(rounds)), rounds, strict=True)
        for d, rounds in num_rounds_by_distances.items()
    )
    parameters_iterator = (
        (
            distance,
            num_rounds,
            noise_model,
            noise_parameter,
            memory_generator,
            noise_parameter_names,
        )
        for ((distance, num_rounds), noise_parameter) in itertools.product(
            distance_and_rounds_iterator, noise_parameters
        )
    )
    # 3. Generate the decoder managers
    decoder_managers: list[StimDecoderManager] = []
    if max_workers == 1:
        for parameters in tqdm(
            parameters_iterator,
            total=total_circuits,
            desc="Generating quantum circuits",
        ):
            decoder_managers.append(
                _generate_surface_code_memory_decoder_manager(*parameters)
            )
        return decoder_managers
    # Else, we might use several workers
    with ProcessPoolExecutor(max_workers=max_workers) as pool:
        return list(
            tqdm(
                iterable=pool.map(
                    _generate_surface_code_memory_decoder_manager_wrapper,
                    parameters_iterator,
                ),
                total=total_circuits,
                desc="Generating quantum circuits",
            )
        )
