import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import roc_curve, auc
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.regularizers import l2
from uad.decision.deep_svdd import anomaly_score_from_predictions, anomaly_score_from_images
from uad.diagnostic.metrics import binarize_set, is_binary
from uad.models.architecture import conv2d_block


class DeepSVDD(Model):
    """
    Support Vector Data Description neural network. Trained on original data and learns a dense embedding while
    trained on the objective function.
    """

    def __init__(self, model, final_dense_shape=32, **kwargs):
        """
        If inputs is None and outputs is None: builds a DeepSVDD network with a LeNet architecture as used in Ruff 2018
        Else give input and outputs to build a model via subclassing
        :param model: either an iterable of layers or a keras.Model
        :param final_dense_shape: # node of the final dense layer
        :param LAMBDA: factor in front of the weight decay loss
        :param kwargs:
        """
        super(DeepSVDD, self).__init__(**kwargs)
        self.CENTER = tf.Variable(initial_value=np.ones(final_dense_shape),
                                  dtype=tf.float32)  # center of the same size as output
        self.RADIUS = 0
        self.model = model

    def set_center(self, new_center):
        self.CENTER = new_center
        tf.print(f"Hypersphere center coordinates: {self.CENTER}")

    def train_step(self, data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            predictions = self(data)
            # distances_to_center = tf.norm(predictions - self.CENTER, axis=-1)
            distances_to_center = tf.math.sqrt(tf.reduce_sum((predictions - self.CENTER) ** 2, axis=1))
            self.RADIUS = tf.reduce_max(distances_to_center)
            centripetal_loss = tf.reduce_mean(distances_to_center, axis=0)
            weight_decay = tf.math.reduce_sum(self.losses)
            total_loss = centripetal_loss + weight_decay
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        return {
            "total_loss": total_loss,
            "centripetal_loss": centripetal_loss,
            "weight_decay": weight_decay
        }

    def test_step(self, data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            predictions = self(data)
            # distances_to_center = tf.norm(predictions - self.CENTER, axis=-1)
            distances_to_center = tf.reduce_sum((predictions - self.CENTER) ** 2, axis=1)
            self.RADIUS = tf.reduce_max(distances_to_center)
            centripetal_loss = tf.reduce_mean(distances_to_center, axis=0)
            weight_decay = tf.math.reduce_sum(self.losses)
            total_loss = centripetal_loss + weight_decay
        return {
            "total_loss": total_loss,
            "centripetal_loss": centripetal_loss,
            "weight_decay": weight_decay
        }

    def call(self, inputs):
        if callable(self.model):
            return self.model(inputs)
        else:
            x = inputs
            for lay in self.model:
                x = lay(x)
            return x

    def score_samples(self, data):
        """
        Returns the anomaly scores for data (name of the method inspired from the sklearn
        interface)
        :param data: image or batch of images
        :return: anomaly scores
        """
        return (anomaly_score_from_images(self, data)).numpy()

    def score_samples_iterator(self, dataset_iterator):
        """
        Compute scores (mse) between images and predictions.
        :param dataset_iterator: iterator of MultiChannelIterator type
        Return: scores in the batch format
        """
        scores = []
        for i in range(len(dataset_iterator)):  # itere a l'infini???
            _, (ims, labs) = dataset_iterator[i]
            if (i + 1) % 50 == 0:
                print(f"making predictions on batch {i + 1}...")
            predictions = self.predict(ims)
            y_scores = tf.math.sqrt(tf.reduce_sum((predictions - ims) ** 2, axis=(1, 2, 3)))
            scores.append(y_scores)

        return np.array(scores)

    def is_anormal(self, data, im_threshold=0):
        predictions = self.predict(data)
        return binarize_set(np.sum((predictions - self.CENTER) ** 2, axis=-1) > im_threshold)

    def compute_AUC(self, fprs, tprs):
        return auc(fprs, tprs)

    def compute_ROC(self, y_true, y_score):
        return roc_curve(y_true, y_score)

    @staticmethod
    def evaluate_on_all(archi, x_train, x_test, y_test, n_classes=10, epochs=30, **params):
        """
        Compute AUC score for the training on each class of the dataset
        """
        auc_scores = []
        for k in range(n_classes):
            print(f"Digit {k}, # Training examples: {x_train[k].shape[0]}")
            model = DeepSVDD(archi, LAMBDA=1e-5)
            model.compile(optimizer=tf.keras.optimizers.Adam())
            model.fit(x_train[k], epochs=epochs, batch_size=128)
            y_score = model.score_samples(x_test)
            fpr, tpr, _ = model.compute_ROC(y_test, y_score)
            auc = model.compute_AUC(fpr, tpr)
            auc_scores.append(auc)

        return np.array(auc_scores)
