from functools import reduce
from typing import cast

import numpy as np

from qibolab._core.components import AcquisitionConfig, Configs, OscillatorConfig
from qibolab._core.execution_parameters import AveragingMode, ExecutionParameters
from qibolab._core.identifier import ChannelId, Result
from qibolab._core.pulses import Acquisition, Readout
from qibolab._core.sequence import PulseSequence
from qibolab._core.sweeper import ParallelSweepers

from .validate import ACQUISITION_MEMORY


def batch_shots(
    sequence: PulseSequence,
    sweepers: list[ParallelSweepers],
    options: ExecutionParameters,
) -> list[int]:
    """Subdivide shots in batches.

    It assumes an integrated/discriminated acquisition, such that each
    acquisition in an individual loop accounts for a single bin.
    """
    assert options.nshots is not None

    if options.averaging_mode is not AveragingMode.SINGLESHOT:
        return [options.nshots]

    bins = np.prod(options.bins(sweepers)[1:])
    acquisitions = max(
        sum(1 for p in pulses if isinstance(p, (Acquisition, Readout)))
        for pulses in sequence.by_channel.values()
    )
    per_shot_memory = bins * acquisitions
    max_shots = int(ACQUISITION_MEMORY // per_shot_memory)
    nfull, remainder = np.divmod(options.nshots, max_shots)
    return [max_shots] * int(nfull) + [int(remainder)]


def concat_shots(
    results: list[dict[int, Result]], options: ExecutionParameters
) -> dict[int, Result]:
    """Concatenate shots batches.

    It collects the measurements' batches resulting from the subdivision generated by
    the :func:`batch_shots` function.
    It assumes homogenuous batches, which are only joined over the single outermost
    dimension, and only for :obj:`AveragingMode.SINGLESHOT` acquisitions.
    """
    return (
        reduce(lambda d, e: {k: np.vstack((v, e[k])) for k, v in d.items()}, results)
        if options.averaging_mode is AveragingMode.SINGLESHOT
        else results[0]
    )


def lo_configs(
    los: dict[ChannelId, str], configs: Configs
) -> dict[ChannelId, OscillatorConfig]:
    """Extract local oscillators' configurations."""
    return {ch: cast(OscillatorConfig, configs[lo]) for ch, lo in los.items()}


def time_of_flights(configs: Configs) -> dict[ChannelId, float]:
    """Extract time of flights configurations."""
    return {
        ch: max(config.delay, 4.0)
        for ch, config in configs.items()
        if isinstance(config, AcquisitionConfig)
    }
