import json
import keras.backend as kb
import numpy as np
import os
import shutil
import warnings
from keras.callbacks import Callback
from sklearn.metrics import roc_auc_score

class MultipleClassAUROC(Callback):
    """
    Monitor mean AUROC and update model
    """
    def __init__(self, sequence, class_names, weights_path, stats=None, workers=1):
        super(Callback, self).__init__()
        self.sequence = sequence
        self.workers = workers
        self.class_names = class_names
        self.weights_path = weights_path
        self.best_weights_path = os.path.join(
            os.path.split(weights_path)[0],"best_"+os.path.split(weights_path)[1],
        )
        self.best_auroc_log_path = os.path.join(
            os.path.split(weights_path)[0],
            "best_auroc.log",
        )
        self.stats_output_path = os.path.join(
            os.path.split(weights_path)[0],
            ".training_stats.json"
        )
        # for resuming previous training
        if stats:
            self.stats = stats
        else:
            self.stats = {"best_mean_auroc": 0}

        # aurocs log
        self.aurocs = {}
        for c in self.class_names:
            self.aurocs[c] = []

    def on_epoch_end(self, epoch, logs={}):
        """
        Calculate the average AUROC and save the best model weights according
        to this metric.

        """
        print("\n*********************************")
        self.stats["lr"] = float(kb.eval(self.model.optimizer.lr))
        print("current learning rate:", self.stats['lr'])

        """
        y_hat shape: (#samples, len(class_names))
        y: [(#samples, 1), (#samples, 1) ... (#samples, 1)]
        """
        y_hat = self.model.predict_generator(self.sequence, workers=self.workers)
        y = self.sequence.get_y_true()

        # print("****************************")
        # print('y_hat', y_hat.shape)
        # print(y_hat)
        # print('y', y.shape)
        # print(y)
        # print("****************************")

        print("*** Epoch# %d dev auroc ***" % (epoch + 1))
        current_auroc = []
        for i in range(len(self.class_names)):
            try:
                score = roc_auc_score(y[:, i], y_hat[:, i])
            except ValueError:
                score = 0
            self.aurocs[self.class_names[i]].append(score)
            current_auroc.append(score)
            print("%d. %s: %f" % ((i+1), self.class_names[i], score))
        print("*********************************")

        # customize your multiple class metrics here
        mean_auroc = np.mean(current_auroc)
        print("mean auroc: %f" % (mean_auroc))
        if mean_auroc > self.stats["best_mean_auroc"]:
            print("Update best auroc from %f to %f" % (self.stats['best_mean_auroc'], mean_auroc))

            # 1. copy best model
            shutil.copy(self.weights_path, self.best_weights_path)

            # # # 2. update log file
            # print("update log file:", self.best_auroc_log_path)
            # with open(self.best_auroc_log_path, "a") as f:
            #     f.write("epoch #%d auroc: %f, lr: %f \n" % ((epoch+1), mean_auroc, self.stats['lr']))

            # 3. write stats output, this is used for resuming the training
            with open(self.stats_output_path, 'w') as f:
                json.dump(self.stats, f)

            print("Update model file: %s -> %s" % (self.weights_path, self.best_weights_path))
            self.stats["best_mean_auroc"] = mean_auroc
            print("*********************************")
        return