from sklearn.metrics import roc_auc_score, precision_recall_curve
from sklearn.metrics import auc as calculate_auc
from sklearn.metrics import mean_squared_error
from sklearn.metrics import accuracy_score
import tensorflow as tf
import os
import numpy as np


from scipy.stats.stats import pearsonr
def r2_score(x,y):
    pcc, _ = pearsonr(x,y)
    return pcc**2

def prc_auc_score(y_true, y_score):
    precision, recall, threshold  = precision_recall_curve(y_true, y_score) #PRC_AUC
    auc = calculate_auc(recall, precision)
    return auc

'''
for early-stopping techniques in regression and classification task
'''
######## Regression ###############################


class Reg_EarlyStoppingAndPerformance(tf.keras.callbacks.Callback):

    def __init__(self, train_data, valid_data, MASK = -1, patience=5, criteria = 'val_loss', verbose = 0):
        super(Reg_EarlyStoppingAndPerformance, self).__init__()
        
        assert criteria in ['val_loss', 'val_r2'], 'not support %s ! only %s' % (criteria, ['val_loss', 'val_r2'])
        self.x, self.y  = train_data
        self.x_val, self.y_val = valid_data
        
        self.history = {'loss':[],
                        'val_loss':[],
                        
                        'rmse':[],
                        'val_rmse':[],
                        
                        'r2':[],
                        'val_r2':[],
                        
                        'epoch':[]}
        self.MASK = MASK
        self.patience = patience
        # best_weights to store the weights at which the minimum loss occurs.
        self.best_weights = None
        self.criteria = criteria
        self.best_epoch = 0
        self.verbose = verbose
        
    def rmse(self, y_true, y_pred):

        N_classes = y_pred.shape[1]
        rmses = []
        for i in range(N_classes):
            y_pred_one_class = y_pred[:,i]
            y_true_one_class = y_true[:, i]
            mask = ~(y_true_one_class == self.MASK)
            mse = mean_squared_error(y_true_one_class[mask], y_pred_one_class[mask])
            rmse = np.sqrt(mse)
            rmses.append(rmse)
        return rmses   
    
    
    def r2(self, y_true, y_pred):
        N_classes = y_pred.shape[1]
        r2s = []
        for i in range(N_classes):
            y_pred_one_class = y_pred[:,i]
            y_true_one_class = y_true[:, i]
            mask = ~(y_true_one_class == self.MASK)
            r2 = r2_score(y_true_one_class[mask], y_pred_one_class[mask])
            r2s.append(r2)
        return r2s   
    
        
    def on_train_begin(self, logs=None):
        # The number of epoch it has waited when loss is no longer minimum.
        self.wait = 0
        # The epoch the training stops at.
        self.stopped_epoch = 0
        # Initialize the best as infinity.
        if self.criteria == 'val_loss':
            self.best = np.Inf  
        else:
            self.best = -np.Inf
            
        
        
 
        
    def on_epoch_end(self, epoch, logs={}):
        
        y_pred = self.model.predict(self.x, verbose=self.verbose)
        rmse_list = self.rmse(self.y, y_pred)
        rmse_mean = np.nanmean(rmse_list)
        
        r2_list = self.r2(self.y, y_pred) 
        r2_mean = np.nanmean(r2_list)
        
        
        y_pred_val = self.model.predict(self.x_val, verbose=self.verbose)
        rmse_list_val = self.rmse(self.y_val, y_pred_val)        
        rmse_mean_val = np.nanmean(rmse_list_val)
        
        r2_list_val = self.r2(self.y_val, y_pred_val)       
        r2_mean_val = np.nanmean(r2_list_val)        
        
        self.history['loss'].append(logs.get('loss'))
        self.history['val_loss'].append(logs.get('val_loss'))
        
        self.history['rmse'].append(rmse_mean)
        self.history['val_rmse'].append(rmse_mean_val)
        
        self.history['r2'].append(r2_mean)
        self.history['val_r2'].append(r2_mean_val)        
        
        self.history['epoch'].append(epoch)
        
        
        # logs is a dictionary
        eph = str(epoch+1).zfill(4)   
        loss = '{0:.4f}'.format((logs.get('loss')))
        val_loss = '{0:.4f}'.format((logs.get('val_loss')))
        rmse = '{0:.4f}'.format(rmse_mean)
        rmse_val = '{0:.4f}'.format(rmse_mean_val)
        r2_mean = '{0:.4f}'.format(r2_mean)
        r2_mean_val = '{0:.4f}'.format(r2_mean_val)
        
        if self.verbose:
            print('\repoch: %s, loss: %s - val_loss: %s; rmse: %s - rmse_val: %s;  r2: %s - r2_val: %s' % (eph,
                                                                                                           loss, val_loss, 
                                                                                                           rmse,rmse_val,
                                                                                                           r2_mean,r2_mean_val),
                  end=100*' '+'\n')


        if self.criteria == 'val_loss':
            current = logs.get(self.criteria)
            if current <= self.best:
                self.best = current
                self.wait = 0
                # Record the best weights if current results is better (less).
                self.best_weights = self.model.get_weights()
                self.best_epoch = epoch

            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.stopped_epoch = epoch
                    self.model.stop_training = True
                    print('\nRestoring model weights from the end of the best epoch.')
                    self.model.set_weights(self.best_weights)    
                    
        else:
            current = np.nanmean(r2_list_val)
            
            if current >= self.best:
                self.best = current
                self.wait = 0
                # Record the best weights if current results is better (less).
                self.best_weights = self.model.get_weights()
                self.best_epoch = epoch

            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.stopped_epoch = epoch
                    self.model.stop_training = True
                    print('\nRestoring model weights from the end of the best epoch.')
                    self.model.set_weights(self.best_weights)              
    
    def on_train_end(self, logs=None):
        self.model.set_weights(self.best_weights)
        if self.stopped_epoch > 0:
            print('\nEpoch %05d: early stopping' % (self.stopped_epoch + 1))

        
        
    def evaluate(self, testX, testY):
        """evalulate, return rmse and r2"""
        y_pred = self.model.predict(testX, verbose=self.verbose)
        rmse_list = self.rmse(testY, y_pred)
        r2_list = self.r2(testY, y_pred)
        return rmse_list, r2_list       



    
    
######## classification ###############################

class CLA_EarlyStoppingAndPerformance(tf.keras.callbacks.Callback):

    def __init__(self, train_data, valid_data, MASK = -1, patience=5, criteria = 'val_loss', metric = 'ROC', last_avf = None, verbose = 0):
        super(CLA_EarlyStoppingAndPerformance, self).__init__()
        
        sp = ['val_loss', 'val_metric']
        assert criteria in sp, 'not support %s ! only %s' % (criteria, sp)
        ms = ['ROC', 'PRC', 'ACC']
        assert metric in ms, 'not support %s ! only %s' % (metric, ms)        
        ms_dict = {'ROC':'roc_auc', 'PRC':'prc_auc', 'ACC':'accuracy'}
        
        metric = ms_dict[metric]
        val_metric = 'val_%s' % metric
        self.metric = metric
        self.val_metric = val_metric
        
        self.x, self.y  = train_data
        self.x_val, self.y_val = valid_data
        self.last_avf = last_avf
        
        
        
        self.history = {'loss':[],
                        'val_loss':[],
                        self.metric:[],
                        self.val_metric:[],
                        'epoch':[]}
        self.MASK = MASK
        self.patience = patience
        # best_weights to store the weights at which the minimum loss occurs.
        self.best_weights = None
        self.criteria = criteria

        
        self.best_epoch = 0
        self.verbose = verbose
        
    def sigmoid(self, x):
        s = 1/(1+np.exp(-x))
        return s

    
    def roc_auc(self, y_true, y_pred):
        if self.last_avf == None:
            y_pred_logits = self.sigmoid(y_pred)
        else:
            y_pred_logits = y_pred
            
        N_classes = y_pred_logits.shape[1]

        aucs = []
        for i in range(N_classes):
            y_pred_one_class = y_pred_logits[:,i]
            y_true_one_class = y_true[:, i]
            mask = ~(y_true_one_class == self.MASK)
            try:
                if self.metric == 'roc_auc':
                    auc = roc_auc_score(y_true_one_class[mask], y_pred_one_class[mask], average='weighted') #ROC_AUC
                elif self.metric == 'prc_auc': 
                    auc = prc_auc_score(y_true_one_class[mask], y_pred_one_class[mask]) #PRC_AUC
                elif self.metric == 'accuracy':
                    auc = accuracy_score(y_true_one_class[mask], np.round(y_pred_one_class[mask])) #ACC
            except:
                auc = np.nan
            aucs.append(auc)
        return aucs  
    
        
        
    def on_train_begin(self, logs=None):
        # The number of epoch it has waited when loss is no longer minimum.
        self.wait = 0
        # The epoch the training stops at.
        self.stopped_epoch = 0
        # Initialize the best as infinity.
        if self.criteria == 'val_loss':
            self.best = np.Inf  
        else:
            self.best = -np.Inf

            
    def on_epoch_end(self, epoch, logs={}):
        
        y_pred = self.model.predict(self.x, verbose = self.verbose)
        roc_list = self.roc_auc(self.y, y_pred)
        roc_mean = np.nanmean(roc_list)
        
        y_pred_val = self.model.predict(self.x_val, verbose = self.verbose)
        roc_val_list = self.roc_auc(self.y_val, y_pred_val)        
        roc_val_mean = np.nanmean(roc_val_list)
        
        self.history['loss'].append(logs.get('loss'))
        self.history['val_loss'].append(logs.get('val_loss'))
        self.history[self.metric].append(roc_mean)
        self.history[self.val_metric].append(roc_val_mean)
        self.history['epoch'].append(epoch)
        
        eph = str(epoch+1).zfill(4)        
        loss = '{0:.4f}'.format((logs.get('loss')))
        val_loss = '{0:.4f}'.format((logs.get('val_loss')))
        auc = '{0:.4f}'.format(roc_mean)
        auc_val = '{0:.4f}'.format(roc_val_mean)    
        
        if self.verbose:
            print('\repoch: %s, loss: %s - val_loss: %s; %s: %s - %s: %s' % (eph,
                                                                             loss, 
                                                                             val_loss, 
                                                                             self.metric,
                                                                             auc,
                                                                             self.val_metric,
                                                                             auc_val), end=100*' '+'\n')

        if self.criteria == 'val_loss':
            current = logs.get(self.criteria)
            if current <= self.best:
                self.best = current
                self.wait = 0
                # Record the best weights if current results is better (less).
                self.best_weights = self.model.get_weights()
                self.best_epoch = epoch

            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.stopped_epoch = epoch
                    self.model.stop_training = True
                    print('\nRestoring model weights from the end of the best epoch.')
                    self.model.set_weights(self.best_weights)    
                    
        else:
            current = roc_val_mean
            if current >= self.best:
                self.best = current
                self.wait = 0
                # Record the best weights if current results is better (less).
                self.best_weights = self.model.get_weights()
                self.best_epoch = epoch

            else:
                self.wait += 1
                if self.wait >= self.patience:
                    self.stopped_epoch = epoch
                    self.model.stop_training = True
                    print('\nRestoring model weights from the end of the best epoch.')
                    self.model.set_weights(self.best_weights)              
    
    def on_train_end(self, logs=None):
        self.model.set_weights(self.best_weights)
        if self.stopped_epoch > 0:
            print('\nEpoch %05d: early stopping' % (self.stopped_epoch + 1))

        
    def evaluate(self, testX, testY):
        
        y_pred = self.model.predict(testX, verbose = self.verbose)
        roc_list = self.roc_auc(testY, y_pred)
        return roc_list