import logging
import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import kld, categorical_crossentropy
from typing import Callable, Dict, Tuple, Union
from alibi_detect.models.tensorflow.trainer import trainer
from alibi_detect.models.tensorflow.losses import loss_distillation
from alibi_detect.utils.prediction import predict_batch
from alibi_detect.base import (BaseDetector, FitMixin, ThresholdMixin,
                               adversarial_prediction_dict)

logger = logging.getLogger(__name__)


class ModelDistillation(BaseDetector, FitMixin, ThresholdMixin):

    def __init__(self,
                 threshold: float = None,
                 distilled_model: tf.keras.Model = None,
                 model: tf.keras.Model = None,
                 loss_type: str = 'kld',
                 temperature: float = 1.,
                 data_type: str = None
                 ) -> None:
        """
        Model distillation concept drift and adversarial detector.

        Parameters
        ----------
        threshold
            Threshold used for score to determine adversarial instances.
        distilled_model
            A tf.keras model to distill.
        model
            A trained tf.keras classification model.
        loss_type
            Loss for distillation. Supported: 'kld', 'xent'
        temperature
            Temperature used for model prediction scaling.
            Temperature <1 sharpens the prediction probability distribution.
        data_type
            Optionally specifiy the data type (tabular, image or time-series). Added to metadata.
        """
        super().__init__()

        if threshold is None:
            logger.warning('No threshold level set. Need to infer threshold using `infer_threshold`.')

        self.threshold = threshold
        self.model = model
        for layer in self.model.layers:  # freeze model layers
            layer.trainable = False

        if isinstance(distilled_model, tf.keras.Model):
            self.distilled_model = distilled_model
        else:
            raise TypeError('No valid format detected for `distilled_model` (tf.keras.Model) ')
        self.loss_type = loss_type
        self.temperature = temperature

        # set metadata
        self.meta['detector_type'] = 'offline'
        self.meta['data_type'] = data_type

    def fit(self,
            X: np.ndarray,
            loss_fn: tf.keras.losses = loss_distillation,
            optimizer: tf.keras.optimizers = tf.keras.optimizers.Adam(learning_rate=1e-3),
            epochs: int = 20,
            batch_size: int = 128,
            verbose: bool = True,
            log_metric: Tuple[str, "tf.keras.metrics"] = None,
            callbacks: tf.keras.callbacks = None,
            preprocess_fn: Callable = None
            ) -> None:
        """
        Train ModelDistillation detector.

        Parameters
        ----------
        X
            Training batch.
        loss_fn
            Loss function used for training.
        optimizer
            Optimizer used for training.
        epochs
            Number of training epochs.
        batch_size
            Batch size used for training.
        verbose
            Whether to print training progress.
        log_metric
            Additional metrics whose progress will be displayed if verbose equals True.
        callbacks
            Callbacks used during training.
        preprocess_fn
            Preprocessing function applied to each training batch.
        """
        # train arguments
        args = [self.distilled_model, loss_fn, X]
        kwargs = {
            'optimizer': optimizer,
            'epochs': epochs,
            'batch_size': batch_size,
            'verbose': verbose,
            'log_metric': log_metric,
            'callbacks': callbacks,
            'preprocess_fn': preprocess_fn,
            'loss_fn_kwargs': {
                'model': self.model,
                'loss_type': self.loss_type,
                'temperature': self.temperature
            }
        }

        # train
        trainer(*args, **kwargs)

    def infer_threshold(self,
                        X: np.ndarray,
                        threshold_perc: float = 99.,
                        margin: float = 0.,
                        batch_size: int = int(1e10)
                        ) -> None:
        """
        Update threshold by a value inferred from the percentage of instances considered to be
        adversarial in a sample of the dataset.

        Parameters
        ----------
        X
            Batch of instances.
        threshold_perc
            Percentage of X considered to be normal based on the adversarial score.
        margin
            Add margin to threshold. Useful if adversarial instances have significantly higher scores and there
            is no adversarial instance in X.
        batch_size
            Batch size used when computing scores.
        """
        # compute adversarial scores
        adv_score = self.score(X, batch_size=batch_size)

        # update threshold
        self.threshold = np.percentile(adv_score, threshold_perc) + margin

    def score(self, X: np.ndarray, batch_size: int = int(1e10), return_predictions: bool = False) \
            -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray, np.ndarray]]:
        """
        Compute adversarial scores.

        Parameters
        ----------
        X
            Batch of instances to analyze.
        batch_size
            Batch size used when computing scores.
        return_predictions
            Whether to return the predictions of the classifier on the original and reconstructed instances.

        Returns
        -------
        Array with adversarial scores for each instance in the batch.
        """

        # model predictions
        y = predict_batch(self.model, X, batch_size=batch_size, proba=True)
        y_distilled = predict_batch(self.distilled_model, X, batch_size=batch_size, proba=True)

        # scale predictions
        if self.temperature != 1.:
            y = y ** (1 / self.temperature)  # type: ignore
            y = (y / tf.reshape(tf.reduce_sum(y, axis=-1), (-1, 1))).numpy()

        if self.loss_type == 'kld':
            score = kld(y, y_distilled).numpy()
        elif self.loss_type == 'xent':
            score = categorical_crossentropy(y, y_distilled).numpy()
        else:
            raise NotImplementedError

        if return_predictions:
            return score, y, y_distilled
        else:
            return score

    def predict(self, X: np.ndarray, batch_size: int = int(1e10), return_instance_score: bool = True) \
            -> Dict[Dict[str, str], Dict[str, np.ndarray]]:
        """
        Predict whether instances are adversarial instances or not.

        Parameters
        ----------
        X
            Batch of instances.
        batch_size
            Batch size used when computing scores.
        return_instance_score
            Whether to return instance level adversarial scores.

        Returns
        -------
        Dictionary containing 'meta' and 'data' dictionaries.
        'meta' has the model's metadata.
        'data' contains the adversarial predictions and instance level adversarial scores.
        """
        score = self.score(X, batch_size=batch_size)

        # values above threshold are adversarial
        pred = (score > self.threshold).astype(int)  # type: ignore

        # populate output dict
        ad = adversarial_prediction_dict()
        ad['meta'] = self.meta
        ad['data']['is_adversarial'] = pred
        if return_instance_score:
            ad['data']['instance_score'] = score
        return ad
