import warnings
from collections.abc import Mapping, Sequence

import numpy as np
import numpy.typing as npt
import pandas as pd

from deltakit_explorer.analysis import (
    LogicalErrorProbabilityPerRoundResults,
    compute_logical_error_per_round,
)
from deltakit_explorer.analysis._analysis import calculate_lep_and_lep_stddev
from deltakit_explorer.analysis._lambda import (
    LambdaResults,
    calculate_lambda_and_lambda_stddev,
)


def _filter_non_close_noise_parameters(
    data: pd.DataFrame,
    noise_parameters: npt.NDArray[np.floating],
    noise_parameter_names: Sequence[str],
) -> pd.DataFrame:
    """Return a filtered view over ``data`` such that all entries are close to the
    provided ``noise_parameters``."""
    ret = data
    for name, param in zip(noise_parameter_names, noise_parameters):
        ret = ret[np.isclose(ret[f"noise_{name}"], param)]
    return ret


def compute_lambda_and_stddev_from_results(
    xi: npt.NDArray[np.floating],
    noise_parameter_names: Sequence[str],
    num_rounds_by_distance: Mapping[int, Sequence[int]],
    data: pd.DataFrame,
) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
    """Compute Λ from ``data`` for all the provided noise parameters in ``xi``.

    This function assumes that the provided ``data`` has been generated with
    ``RunAllAnalysisEngine`` and the tasks returned by
    ``generate_decoder_managers_for_lambda`` called on the exact same parameters as
    provided here.

    Args:
        xi (npt.NDArray[numpy.floating]): an array of shape ``(m, n)``, where ``m`` is
            the number of noise parameters needed by ``noise_model_type`` and ``n`` the
            number of points to evaluate Λ on.
        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.
        data (pandas.DataFrame): data generated by executing ``RunAllAnalysisEngine`` on
            the tasks returned by ``generate_decoder_managers_for_lambda`` called
            with the exact same parameters as provided to this function.

    Returns:
        two arrays of dimensions ``(1, n)``. The first array contains an estimation of Λ
        for each parameters in ``xi`` and the second array contains the computed
        standard deviations.
    """
    if len(xi.shape) != 2:
        msg = f"Expected a 2-dimensional array but got shape {xi.shape}."
        raise ValueError(msg)
    _, n = xi.shape
    ret: npt.NDArray[np.floating] = np.zeros((1, n), dtype=np.float64)
    stddev: npt.NDArray[np.floating] = np.zeros_like(ret)
    for i in range(n):
        filtered_data = _filter_non_close_noise_parameters(
            data, xi[:, i], noise_parameter_names
        )
        res = _compute_lambda_from_results(num_rounds_by_distance, filtered_data)
        ret[0, i], stddev[0, i] = res.lambda_, res.lambda_stddev
    return ret, stddev


def _compute_lambda_from_results(
    num_rounds_by_distance: Mapping[int, Sequence[int]],
    data: pd.DataFrame,
) -> LambdaResults:
    lepprs: list[float] = []
    leppr_stddevs: list[float] = []
    distances = sorted(num_rounds_by_distance.keys())
    for d in distances:
        leppr = _compute_logical_error_rate_per_round_from_results(
            num_rounds_by_distance[d], data[data["distance"] == d]
        )
        lepprs.append(leppr.leppr)
        leppr_stddevs.append(leppr.leppr_stddev)
    return calculate_lambda_and_lambda_stddev(distances, lepprs, leppr_stddevs)


def _compute_logical_error_rate_per_round_from_results(
    num_rounds: npt.NDArray[np.int_] | Sequence[int], data: pd.DataFrame
) -> LogicalErrorProbabilityPerRoundResults:
    num_fails: npt.NDArray[np.int_] | list[int] = []
    max_shots: npt.NDArray[np.int_] | list[int] = []
    for nrounds in num_rounds:
        data_row = data.query(f"num_rounds == {nrounds}")
        nfails = data_row["fails"].to_numpy()[0]
        nshots = data_row["shots"].to_numpy()[0]
        num_fails.append(nfails)
        max_shots.append(nshots)
    # Filter out 0 fails
    non_zeros_mask = np.asarray(num_fails) != 0
    if not np.all(non_zeros_mask):
        warnings.warn(
            "Found at least one sampling task with 0 fails. It will be ignored for "
            "logical error probability computation."
        )
    num_rounds = np.asarray(num_rounds)[non_zeros_mask]
    num_fails = np.asarray(num_fails)[non_zeros_mask]
    max_shots = np.asarray(max_shots)[non_zeros_mask]

    lep, lep_stddev = calculate_lep_and_lep_stddev(num_fails, max_shots)
    return compute_logical_error_per_round(num_rounds, lep, lep_stddev)
