"""
Module related to abstract explainer
"""

from abc import ABC, abstractmethod
import warnings

import tensorflow as tf

from ..types import Callable, Dict, Tuple, Union, Optional
from ..utils import find_layer


class BlackBoxExplainer(ABC):
    """
    Base class for Black-Box explainers.

    Parameters
    ----------
    model
        Model used for computing explanations.
    batch_size
        Number of samples to explain at once, if None compute all at once.
    """

    # in order to avoid re-tracing at each tf.function call,
    # share the reconfigured models between the methods if possible
    _cache_models: Dict[Tuple[int, int], tf.keras.Model] = {}

    def __init__(self, model: Callable, batch_size: Optional[int] = 64):
        if isinstance(model, tf.keras.Model):
            model_key = (id(model.input), id(model.output))
            if model_key not in BlackBoxExplainer._cache_models:
                BlackBoxExplainer._cache_models[model_key] = model
            self.model = BlackBoxExplainer._cache_models[model_key]
        else:
            self.model = model

        self.batch_size = batch_size

    @abstractmethod
    def explain(self,
                inputs: tf.Tensor,
                labels: tf.Tensor) -> tf.Tensor:
        """
        Compute the explanations of the given samples.

        Parameters
        ----------
        inputs
            Input samples to be explained.
        labels
            One-hot encoded labels, one for each sample.

        Returns
        -------
        explanations
            Explanation generated by the method.
        """
        raise NotImplementedError()

    def __call__(self,
                 inputs: tf.Tensor,
                 labels: tf.Tensor) -> tf.Tensor:
        """Explain alias"""
        return self.explain(inputs, labels)

    @staticmethod
    @tf.function
    def _predictions(model: Callable,
                     inputs: tf.Tensor,
                     labels: tf.Tensor) -> tf.Tensor:
        """
        Compute predictions scores, only for the label class, for a batch of samples.

        Parameters
        ----------
        model
            Model used for computing predictions.
        inputs
            Input samples to be explained.
        labels
            One-hot encoded labels, one for each sample.

        Returns
        -------
        scores
            Predictions scores computed, only for the label class.
        """
        scores = tf.reduce_sum(model(inputs) * labels, axis=-1)
        return scores

    @staticmethod
    def _batch_predictions(model: Callable,
                           inputs: tf.Tensor,
                           labels: tf.Tensor,
                           batch_size: Optional[int]) -> tf.Tensor:
        """
        Compute predictions scores, only for the label class, for the samples passed. Take
        care of splitting in multiple batches if batch_size is specified.

        Parameters
        ----------
        model
            Model used for computing predictions score.
        inputs
            Input samples to be explained.
        labels
            One-hot encoded labels, one for each sample.
        batch_size
            Number of samples to predict at once, if None compute all at once.

        Returns
        -------
        scores
            Predictions scores computed, only for the label class.
        """
        if batch_size is not None:
            dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))
            scores = tf.concat([
                BlackBoxExplainer._predictions(model, x_batch, y_batch)
                for x_batch, y_batch in dataset.batch(batch_size)
            ], axis=0)
        else:
            scores = BlackBoxExplainer._predictions(model, inputs, labels)

        return scores


class WhiteBoxExplainer(BlackBoxExplainer, ABC):
    """
    Base class for White-Box explainers.

    Parameters
    ----------
    model
        Model used for computing explanations.
    output_layer
        Layer to target for the output (e.g logits or after softmax), if int, will be be interpreted
        as layer index, if string will look for the layer name. Default to the last layer, it is
        recommended to use the layer before Softmax.
    batch_size
        Number of samples to explain at once, if None compute all at once.
    """

    def __init__(self,
                 model: tf.keras.Model,
                 output_layer: Optional[Union[str, int]] = None,
                 batch_size: Optional[int] = 64):

        if output_layer is not None:
            # reconfigure the model (e.g skip softmax to target logits)
            target_layer = find_layer(model, output_layer)
            model = tf.keras.Model(model.input, target_layer.output)

            # sanity check, output layer before softmax
            try:
                if target_layer.activation.__name__ == tf.keras.activations.softmax.__name__:
                    warnings.warn("Output is after softmax, it is recommended to "
                                  "use the layer before.")
            except AttributeError:
                pass

        super().__init__(model, batch_size)

    @staticmethod
    @tf.function
    def _gradient(model: Callable,
                  inputs: tf.Tensor,
                  labels: tf.Tensor) -> tf.Tensor:
        """
        Compute gradients for a batch of samples.

        Parameters
        ----------
        model
            Model used for computing gradient.
        inputs
            Input samples to be explained.
        labels
            One-hot encoded labels, one for each sample.

        Returns
        -------
        gradients
            Gradients computed, with the same shape as the inputs.
        """
        with tf.GradientTape(watch_accessed_variables=False) as tape: # type: ignore
            tape.watch(inputs)
            score = tf.reduce_sum(tf.multiply(model(inputs), labels), axis=1)
        return tape.gradient(score, inputs)

    @staticmethod
    def _batch_gradient(model: Callable,
                        inputs: tf.Tensor,
                        labels: tf.Tensor,
                        batch_size: Optional[int]) -> tf.Tensor:
        """
        Compute the gradients of the sample passed, take care of splitting the samples in
        multiple batches if batch_size is specified.

        Parameters
        ----------
        model
            Model used for computing gradient.
        inputs
            Input samples to be explained.
        labels
            One-hot encoded labels, one for each sample.
        batch_size
            Number of samples to explain at once, if None compute all at once.

        Returns
        -------
        gradients
            Gradients computed, with the same shape as the inputs.
        """
        if batch_size is not None:
            dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))
            gradients = tf.concat([
                WhiteBoxExplainer._gradient(model, x, y)
                for x, y in dataset.batch(batch_size)
            ], axis=0)
        else:
            gradients = WhiteBoxExplainer._gradient(model, inputs, labels)

        return gradients
