"""
Module related to abstract explainer
"""

from abc import ABC, abstractmethod
import warnings

import tensorflow as tf
import numpy as np

from ..types import Callable, Dict, Tuple, Union, Optional
from ..commons import find_layer, tensor_sanitize


def sanitize_input_output(explanation_method: Callable):
    """
    Wrap a method explanation function to ensure tf.Tensor as inputs,
    and as output

    explanation_method
        Function to wrap, should return an tf.tensor.
    """
    def sanitize(self, inputs: Union[tf.data.Dataset, tf.Tensor, np.array],
                 targets: Optional[Union[tf.Tensor, np.array]],
                 *args):
        # ensure we have tf.tensor
        inputs, targets = tensor_sanitize(inputs, targets)
        # then enter the explanation function
        return explanation_method(self, inputs, targets, *args)

    return sanitize


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: Union[tf.data.Dataset, tf.Tensor, np.array],
                targets: Optional[Union[tf.Tensor, np.array]] = None) -> tf.Tensor:
        """
        Compute the explanations of the given samples.
        Accept Tensor, numpy array or tf.data.Dataset (in that case targets is None)

        Parameters
        ----------
        inputs
            Input samples to be explained.
        targets
            One-hot encoded labels or regression target (e.g {+1, -1}), 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)


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)
