from typing import Optional, Any
from torch import Tensor
from optuna.trial import Trial, TrialState

from pnpxai.core._types import TensorOrTupleOfTensors
from pnpxai.core.modality.modality import Modality, TextModality
from pnpxai.explainers import Explainer, KernelShap, Lime
from pnpxai.explainers.utils.postprocess import PostProcessor, Identity
from pnpxai.evaluator.metrics.base import Metric
from pnpxai.evaluator.optimizer.suggestor import suggest
from pnpxai.utils import format_into_tuple, format_out_tuple_if_single, generate_param_key


class Objective:
    """
    A class that encapsulates the logic for evaluating a model's performance using a 
    specified explainer, postprocessor, and metric within a given modality. The `Objective` 
    class is designed to be callable and can be used within an optimization framework 
    like Optuna to evaluate different configurations of explainers and postprocessors.

    Parameters:
        explainer (Explainer): 
            The explainer used to generate attributions for the model.
        postprocessor (PostProcessor): 
            The postprocessor applied to the attributions generated by the explainer.
        metric (Metric): 
            The metric used to evaluate the effectiveness of the postprocessed attributions.
        modality (Modality): 
            The data modality (e.g., image, text) the model and explainer are operating on.
        inputs (Optional[TensorOrTupleOfTensors], optional): 
            The input data to the model. Defaults to None.
        targets (Optional[Tensor], optional): 
            The target labels for the input data. Defaults to None.
    """

    EXPLAINER_KEY = 'explainer'
    POSTPROCESSOR_KEY = 'postprocessor'

    def __init__(
        self,
        explainer: Explainer,
        postprocessor: PostProcessor,
        metric: Metric,
        modality: Modality,
        inputs: Optional[TensorOrTupleOfTensors] = None,
        targets: Optional[Tensor] = None,
    ):
        self.explainer = explainer
        self.postprocessor = postprocessor
        self.metric = metric
        self.modality = modality
        self.inputs = inputs
        self.targets = targets

    def set_inputs(self, inputs):
        """
        Sets the input data for the objective.

        Parameters:
            inputs (TensorOrTupleOfTensors): 
                The input data to be used by the explainer.

        Returns:
            Objective: The updated Objective instance.
        """
        self.inputs = inputs
        return self

    def set_targets(self, targets):
        """
        Sets the target labels for the objective.

        Parameters:
            targets (Tensor): 
                The target labels corresponding to the input data.

        Returns:
            Objective: The updated Objective instance.
        """
        self.targets = targets
        return self

    def set_data(self, inputs, targets):
        """
        Sets both the input data and target labels for the objective.

        Parameters:
            inputs (TensorOrTupleOfTensors): 
                The input data to be used by the explainer.
            targets (Tensor): 
                The target labels corresponding to the input data.

        Returns:
            Objective: The updated Objective instance.
        """
        self.set_inputs(inputs)
        self.set_targets(targets)
        return self

    def __call__(self, trial: Trial) -> float:
        """
        Executes the objective function within an optimization trial. This method is 
        responsible for selecting the explainer and postprocessor based on the trial 
        suggestions, performing the explanation and postprocessing, and evaluating the 
        results using the specified metric.

        Parameters:
            trial (Trial): 
                The trial object from an optimization framework like Optuna.

        Returns:
            float: The evaluation result of the metric after applying the explainer and 
            postprocessor to the model's predictions. Returns `nan` if the postprocessed 
            results contain non-countable values like `nan` or `inf`.
        """
        # suggest explainer
        explainer = suggest(trial, self.explainer, self.modality, key=self.EXPLAINER_KEY)

        # suggest postprocessor
        modalities = format_into_tuple(self.modality)
        is_multi_modal = len(modalities) > 1
        postprocessor = []
        for pp, modality in zip(
            format_into_tuple(self.postprocessor),
            modalities,
        ):
            force_params = {}
            if (
                isinstance(explainer, (Lime, KernelShap))
                and isinstance(modality, TextModality)
            ):
                force_params['pooling_fn'] = Identity()
            postprocessor.append(suggest(
                trial, pp, modality,
                key=generate_param_key(
                    self.POSTPROCESSOR_KEY,
                    modality.__class__.__name__ if is_multi_modal else None
                ),
                force_params=force_params,
            ))
        postprocessor = format_out_tuple_if_single(tuple(postprocessor))

        # Ignore duplicated samples
        states_to_consider = (TrialState.COMPLETE,)
        trials_to_consider = trial.study.get_trials(deepcopy=False, states=states_to_consider)
        for t in reversed(trials_to_consider):
            if trial.params == t.params:
                trial.set_user_attr('explainer', explainer)
                trial.set_user_attr('postprocessor', postprocessor)
                return t.value

        # Explain and postprocess
        attrs = explainer.attribute(self.inputs, self.targets)
        postprocessed = tuple(
            pp(attr) for pp, attr in zip(
                format_into_tuple(postprocessor),
                format_into_tuple(attrs),
            )
        )

        if any(pp.isnan().sum() > 0 or pp.isinf().sum() > 0 for pp in postprocessed):
            # Treat a failure as nan
            return float('nan')

        postprocessed = format_out_tuple_if_single(postprocessed)
        metric = self.metric.set_explainer(explainer)
        evals = format_into_tuple(
            metric.evaluate(self.inputs, self.targets, postprocessed)
        )

        # Keep current explainer and postprocessor on trial
        trial.set_user_attr('explainer', explainer)
        trial.set_user_attr('postprocessor', postprocessor)
        return (sum(*evals) / len(evals)).item()
