from typing import (
    Callable,
    List,
    Optional,
    Tuple,
    Union,
)

from jax import vmap
import jax.numpy as jnp
import numpy as np

from fortuna.conformal.batch_mvp import BatchMVPConformalMethod
from fortuna.conformal.classification.base import ConformalClassifier
from fortuna.data.loader import (
    DataLoader,
    InputsLoader,
)
from fortuna.typing import Array


class BatchMVPConformalClassifier(BatchMVPConformalMethod, ConformalClassifier):
    def __init__(
        self,
        score_fn: Callable[[Array, Array], Array],
        group_fns: List[Callable[[Array], Array]],
        n_classes: int,
        n_buckets: int = 100,
    ):
        """
        This class implements a classification version of BatchMVP
        `[Jung et al., 2022] <https://arxiv.org/abs/2209.15145>`_,
        a multivalid conformal prediction method that satisfies coverage guarantees conditioned on group membership
        and non-conformity threshold.

        Parameters
        ----------
        score_fn: Callable[[Array, Array], Array]
            A score function mapping a batch of inputs and targets to scalar scores, one for each data point. The
            score function represents the degree of non-conformity between inputs and targets. In regression, an
            example of score function is :math:`s(x,y)=|y - h(x)|`, where `h` is an arbitrary regression model.
        group_fns: List[Callable[[Array], Array]]
            A list of group functions, each mapping input data points into boolean arrays which determine whether
            an input belongs to a certain group or not. As an example, suppose that we are interested in obtaining
            marginal coverages guarantee on both negative and positive scalar inputs.
            Then we could define groups functions
            :math:`g_1(x) = x < 0` and :math:`g_1(x) = x > 0`.
            Note that groups can be overlapping, and do not need to cover the full space of inputs.
        n_classes: int
            The number of distinct classes to classify among. The underlying assumption is that the classes are
            identified with an integer from 0 to :code:`n_classes-1`.
        n_buckets: int
            The number of buckets that defines the search space between 0 and 1 that determines the updates of the
            thresholds for the score function.
        """
        super().__init__(score_fn=score_fn, group_fns=group_fns, n_buckets=n_buckets)
        self.n_classes = n_classes

    def conformal_set(
        self,
        val_data_loader: DataLoader,
        test_inputs_loader: InputsLoader,
        error: float = 0.05,
        tol: float = 1e-4,
        n_rounds: int = 1000,
        return_max_calib_error: bool = False,
        test_thresholds: Optional[Array] = None,
    ) -> Union[List[List[int]], Tuple[List[List[int]], List[Array]]]:
        """
        Compute a conformal set for each test input.

        Parameters
        ----------
        val_data_loader: DataLoader
            A data loader of validation data.
        test_inputs_loader: InputsLoader
            A loader of test input data points.
        error: float
            A desired coverage error.
        tol: float
            A tolerance for the maximum calibration error.
        n_rounds: int
            The maximum number of updates the algorithm will run for.
        return_max_calib_error: bool
            Whether to return a list of computed maximum calibration error, that is the larger calibration error
            over the different groups.
        test_thresholds: Optional[Array]
            The score thresholds computed over the test data set. These should be the output of
            `BatchMVP.threshold_score`. If provided, they will not be recomputed internally.

        Returns
        -------
        Union[List[List[int]], Tuple[List[List[int]], List[Array]]]
            The computed conformal sets for each test input. Optionally, it returns the maximum calibration errors
            computed during the algorithm.
        """
        if test_thresholds is not None and return_max_calib_error:
            raise ValueError(
                "If `test_thresholds` is given, `return_max_calib_error` cannot be returned."
            )
        if test_thresholds is None:
            outs = self.threshold_score(
                val_data_loader=val_data_loader,
                test_inputs_loader=test_inputs_loader,
                error=error,
                tol=tol,
                n_rounds=n_rounds,
                return_max_calib_error=return_max_calib_error,
            )
            if return_max_calib_error:
                test_thresholds, max_calib_errors = outs
            else:
                test_thresholds = outs

        c = 0
        all_ys = jnp.arange(self.n_classes)
        all_bools = []
        for inputs in test_inputs_loader:
            batch_thresholds = test_thresholds[c : c + inputs.shape[0]]
            all_bools.append(
                vmap(
                    lambda y: self.score_fn(inputs, y) <= batch_thresholds, out_axes=1
                )(all_ys)
            )
            c += inputs.shape[0]
        all_bools = jnp.concatenate(all_bools, axis=0)

        sizes = np.sum(all_bools, 1)
        sets = np.zeros(c, dtype=object)
        for s in np.unique(sizes):
            idx = jnp.where(sizes == s)[0]
            sets[idx] = np.nonzero(all_bools[idx])[1].reshape(len(idx), s).tolist()
        sets = sets.tolist()

        if return_max_calib_error:
            return sets, max_calib_errors
        return sets
