import multiprocessing
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path

import sinter

from tqec.compile.compile import _DEFAULT_BLOCK_REPETITIONS, compile_block_graph
from tqec.compile.convention import FIXED_BULK_CONVENTION, Convention
from tqec.compile.detectors.database import DetectorDatabase
from tqec.computation.block_graph import BlockGraph
from tqec.computation.correlation import CorrelationSurface
from tqec.simulation.generation import generate_sinter_tasks
from tqec.simulation.split import (
    heuristic_custom_error_key,
    split_stats_for_observables,
)
from tqec.utils.noise_model import NoiseModel
from tqec.utils.paths import DEFAULT_DETECTOR_DATABASE_PATH
from tqec.utils.scale import LinearFunction


def start_simulation_using_sinter(
    block_graph: BlockGraph,
    ks: Sequence[int],
    ps: Sequence[float],
    noise_model_factory: Callable[[float], NoiseModel],
    manhattan_radius: int,
    convention: Convention = FIXED_BULK_CONVENTION,
    observables: list[CorrelationSurface] | None = None,
    detector_database: DetectorDatabase | None = None,
    database_path: str | Path = DEFAULT_DETECTOR_DATABASE_PATH,
    num_workers: int = multiprocessing.cpu_count(),
    progress_callback: Callable[[sinter.Progress], None] | None = None,
    max_shots: int | None = None,
    max_errors: int | None = None,
    decoders: Iterable[str] = ("pymatching",),
    print_progress: bool = False,
    custom_decoders: dict[str, sinter.Decoder | sinter.Sampler] | None = None,
    save_resume_filepath: str | Path | None = None,
    existing_data_filepaths: Iterable[str | Path] = (),
    split_observable_stats: bool = True,
    block_temporal_height: LinearFunction = _DEFAULT_BLOCK_REPETITIONS,
) -> list[list[sinter.TaskStats]]:
    """Run `stim` simulations using `sinter`.

    This function is the preferred entry-point to run `sinter` computations using
    `stim` from circuits generated by `tqec`.

    It removes the need to generate the `stim.Circuit` instances in parallel (due
    to the relative inefficiency of detector annotation at the moment) by
    importing and calling :func:`generate_sinter_tasks`. It also forwards
    several parameters to :func:`sinter.collect`, but without showing in its
    signature arguments that we do not envision using.

    Args:
        block_graph: a representation of the QEC computation to simulate.
        ks: values of the scaling parameter `k` to use in order to generate the
            circuits.
        ps: values of the noise parameter `p` to use to instantiate a noise
            model using the provided `noise_model_factory`.
        noise_model_factory: a callable that is used to instantiate a noise
            model from each of the noise parameters in `ps`.
        manhattan_radius: radius used to automatically compute detectors. The
            best value to set this argument to is the minimum integer such that
            flows generated from any given reset/measurement, from any plaquette
            at any spatial/temporal place in the QEC computation, do not
            propagate outside of the qubits used by plaquettes spatially located
            at maximum `manhattan_radius` plaquettes from the plaquette the
            reset/measurement belongs to (w.r.t. the Manhattan distance).
            Default to 2, which is sufficient for regular surface code. If
            negative, detectors are not computed automatically and are not added
            to the generated circuits.
        convention: convention used to generate the quantum circuits.
        observables: a list of correlation surfaces to compile to logical
             observables and generate statistics for. If `None`, all the correlation
             surfaces of the provided computation are used.
        detector_database: an instance to retrieve from / store in detectors
            that are computed as part of the circuit generation.
        database_path: str or Path where detector database is presaved
        num_workers: The number of worker processes to use.
        progress_callback: Defaults to None (unused). If specified, then each
            time new sample statistics are acquired from a worker this method
            will be invoked with the new `sinter.TaskStats`.
        max_shots: Defaults to None (unused). Stops the sampling process
            after this many samples have been taken from the circuit.
        max_errors: Defaults to None (unused). Stops the sampling process
            after this many errors have been seen in samples taken from the
            circuit. The actual number sampled errors may be larger due to
            batching.
        decoders: Defaults to None (specified by each Task). The names of the
            decoders to use on each Task. It must either be the case that each
            Task specifies a decoder and this is set to None, or this is an
            iterable and each Task has its decoder set to None.
        print_progress: When True, progress is printed to stderr while
            collection runs.
        custom_decoders: Named child classes of `sinter.decoder`, that can be
            used if requested by name by a task or by the decoders list.
            If not specified, only decoders with support built into sinter, such
            as 'pymatching' and 'fusion_blossom', can be used.
        save_resume_filepath: Defaults to None (unused). If set to a filepath,
            results will be saved to that file while they are collected. If the
            python interpreter is stopped or killed, calling this method again
            with the same save_resume_filepath will load the previous results
            from the file so it can resume where it left off. The stats in this
            file will be counted in addition to each task's previous_stats field
            (as opposed to overriding the field). Notes that the stats for each
            individual observable are not split but saved in the `custom_counts`
            field of the stats.
        existing_data_filepaths: CSV data saved to these files will be loaded,
            included in the returned results, and count towards things like
            max_shots and max_errors.
        split_observable_stats: Defaults to True. If True, the results are
            post-processed to get individual statistics for each observable in
            `observables`. If False, the results are returned as they are
            collected.
        block_temporal_height: the number of rounds of stabilizer measurements
            (ignoring one layer for initialization and another for final measurement).
            Defaults to `2k-1`.

    Returns:
        A list of lists of `sinter.TaskStats`. If `split_observable_stats` is
        True, the outer list has one element per provided observable in `observables`.
        If `split_observable_stats` is False, the outer list has only one element,
        containing the raw statistics collected.

    """
    if observables is None:
        observables = block_graph.find_correlation_surfaces()
    custom_error_count_key: str | None = None
    if split_observable_stats and len(observables) > 1:
        custom_error_count_key = heuristic_custom_error_key(observables)

    compiled_graph = compile_block_graph(
        block_graph,
        convention,
        observables,
        block_temporal_height,
    )
    stats = sinter.collect(
        num_workers=num_workers,
        tasks=generate_sinter_tasks(
            compiled_graph,
            ks,
            ps,
            noise_model_factory,
            manhattan_radius,
            detector_database,
            database_path,
        ),
        existing_data_filepaths=existing_data_filepaths,
        save_resume_filepath=save_resume_filepath,
        progress_callback=progress_callback,
        max_shots=max_shots,
        max_errors=max_errors,
        decoders=decoders,
        print_progress=print_progress,
        custom_decoders=custom_decoders,
        hint_num_tasks=len(ks) * len(ps),
        count_observable_error_combos=True,
        custom_error_count_key=custom_error_count_key,
    )
    if split_observable_stats:
        return split_stats_for_observables(stats, len(observables))
    return [stats]
