###############################################################################
# Copyright (c) 2018, Lawrence Livermore National Security, LLC.
#
# Produced at the Lawrence Livermore National Laboratory
#
# Written by K. Humbird (humbird1@llnl.gov), L. Peterson (peterson76@llnl.gov).
#
# LLNL-CODE-754815
#
# All rights reserved.
#
# This file is part of DJINN.
#
# For details, see github.com/LLNL/djinn.
#
# For details about use and distribution, please read DJINN/LICENSE .
###############################################################################

try:
    import tensorflow.compat.v1 as tf
    tf.disable_v2_behavior()
except:
    import tensorflow as tf
import numpy as np
try:
    import cPickle
except:
    import _pickle as cPickle
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import MinMaxScaler
from djinn_fns import tree_to_nn_weights, tf_dropout_regression, \
               get_hyperparams, tf_continue_training



def load(model_name, model_path="./"):
    """Reload model and launch tensorflow session.

    Args:
        model_name (str): Name of model.
        model_path (str): Path to model.

    Returns:
        Object: djinn regressor model.
    """
    try:
        with open("%s%s.pkl"%(model_path,model_name), "rb") as f:
            model=cPickle.load(f)
        model.load_model(model_name, model_path)
        return(model)
    except:
        print("Error loading model.")
        return()


class DJINN_Regressor():
    """DJINN regression model.

    Args:
        n_trees (int): Number of trees in random forest
                (equal to the number of neural networks).
        max_tree_depth (int): Maximum depth of decision tree.
                       Neural network will have max_tree_depth-1
                       hidden layers.
        dropout_keep_prob (float): Probability of keeping a neuron
                          in dropout layers.
    """

    def __init__(self, n_trees=1, max_tree_depth=4, dropout_keep_prob=1.0):
        self.__n_trees = n_trees
        self.__tree_max_depth = max_tree_depth
        self.__dropout_keep_prob = dropout_keep_prob
        self.__yscale = None
        self.__xscale = None
        self.__sess = None
        self.__regression = True
        self.modelname = None
        self.modelpath = None




    def get_hyperparameters(self, X, Y, weight_reg=1.0e-8, random_state=None):
        """Automatic selection of djinn hyper-parameters.

        Returns learning rate, number of epochs, batch size.

        Args:
            X (ndarray): Input parameters for training data.
            Y (ndarray): Target parameters for training data.
            weight_reg (float): Multiplier for L2 penalty on weights.
            random_state (int): Set the random seed.

        Raises:
            Exception: if decision tree cannot be built from the data.

        Returns:
            dictionary: Dictionary with batch size,
                        learning rate, number of epochs
        """

        if (X.ndim == 1):
            print('Please reshape single-input data to a one-column array')
            return


        single_output = False
        if (Y.ndim == 1):
            single_output = True
            Y = Y.reshape(-1,1)

        # Scale the data
        self.__xscale = MinMaxScaler().fit(X)
        if self.__regression == True:
            self.__yscale = MinMaxScaler().fit(Y)


        # Train the random forest
        rfr = RandomForestRegressor(self.__n_trees, max_depth=self.__tree_max_depth,
                                    bootstrap=False, random_state=random_state)
        if self.__regression == True:
            if (single_output == True):
                rfr.fit(self.__xscale.transform(X), self.__yscale.transform(Y).flatten())
            else:
                rfr.fit(self.__xscale.transform(X), self.__yscale.transform(Y))

        else:
            if (single_output == True):
                rfr.fit(self.__xscale.transform(X), Y.flatten())
            else:
                rfr.fit(self.__xscale.transform(X), Y)

        if(rfr.estimators_[0].tree_.max_depth < 1):
            raise Exception("Error: Cannot build decision tree.")

        # Map trees to initialized neural networks
        tree_to_network = tree_to_nn_weights(self.__regression, X, Y, self.__n_trees, rfr, random_state)

        print('Finding optimal hyper-parameters...')
        # Run auto-djinn
        nn_batch_size, learnrate, nn_epochs = get_hyperparams(self.__regression,
                tree_to_network, self.__xscale, self.__yscale, X, Y,
                self.__dropout_keep_prob, weight_reg, random_state=random_state)

        return({'batch_size':nn_batch_size, 'learn_rate':learnrate, 'epochs':nn_epochs})



    def train(self, X, Y, epochs=1000, learn_rate=0.001, batch_size=0, weight_reg=1.0e-8,
              display_step=1, save_files=True, file_name="djinn", save_model=True,
              model_name="djinn_model", model_path="./", random_state=None):
        """Train djinn with specified hyperparameters.

        Args:
            X (ndarray): Input parameters for training data.
            Y (ndarray): Target parameters for training data.
            epochs (int): Number of training epochs.
            learn_rate (float): Learning rate for optimizaiton of weights/biases.
            batch_size (int): Number of samples per batch.
            weight_reg (float): Multiplier for L2 penalty on weights.
            display_step (int): Cost is printed every display_steps during training.
            save_files (bool): If True, saves train/valid cost per epoch, weights/biases.
            file_name (str): File name used if 'save_files' is True.
            save_model (bool): If True, saves the trained model.
            model_name (str): File name for model if 'save_model' is True.
            model_path (str): Location of where the model/files are saved.
            random_state (int): Set the random seed.

        Raises:
            Exception: if decision tree cannot be built from the data.

        Returns:
            None
        """
        self.modelname = model_name
        self.modelpath = model_path

        if (X.ndim == 1):
            print('Please reshape single-input data to a one-column array')
            return

        # Reshape data to play well with sklearn
        single_output = False
        if (Y.ndim == 1):
            single_output = True
            Y = Y.reshape(-1,1)

        # Create scalers
        if(self.__xscale == None):
            self.__xscale = MinMaxScaler().fit(X)
            if self.__regression == True:
                self.__yscale = MinMaxScaler().fit(Y)

        # Train the random forest
        rfr = RandomForestRegressor(self.__n_trees, max_depth=self.__tree_max_depth,
                                    bootstrap=False, random_state=random_state)
        if self.__regression == True:
            if (single_output == True):
                rfr.fit(self.__xscale.transform(X), self.__yscale.transform(Y).flatten())
            else:
                rfr.fit(self.__xscale.transform(X), self.__yscale.transform(Y))

        else:
            if (single_output == True):
                rfr.fit(self.__xscale.transform(X), Y.flatten())
            else:
                rfr.fit(self.__xscale.transform(X), Y)

        # Check the forest was successful
        if(rfr.estimators_[0].tree_.max_depth <= 1):
            raise Exception("Error: Cannot build decision tree.")
        # Map trees to neural networks
        tree_to_network = tree_to_nn_weights(self.__regression, X, Y, self.__n_trees, rfr, random_state)
        if(batch_size == 0): batch_size = int(np.ceil(0.05*len(Y)))
        tf_dropout_regression(self.__regression, tree_to_network, self.__xscale,
                             self.__yscale, X, Y,ntrees=self.__n_trees,
                             filename=file_name, learnrate=learn_rate,
                             training_epochs=epochs, batch_size=batch_size,
                             dropout_keep_prob=self.__dropout_keep_prob, weight_reg=weight_reg,
                             display_step=display_step, savefiles=save_files,
                             savemodel=save_model, modelname=self.modelname,
                             modelpath=self.modelpath, random_state=random_state)
        if (save_model == True):
            with open('%s%s.pkl'%(self.modelpath, self.modelname), 'wb') as f:
                cPickle.dump(self, f)




    def fit(self, X, Y, epochs=None, learn_rate=None, batch_size=None, weight_reg=1.0e-8,
            display_step=1, save_files=True, file_name="djinn", save_model=True,
            model_name="djinn_model", model_path="./", random_state=None):
        """Trains djinn model with optimal settings, if not supplied.

        Args:
            X (ndarray): Input parameters for training data.
            Y (ndarray): Target parameters for training data.
            epochs (int): Number of training epochs.
            learn_rate (float): Learning rate for optimizaiton of weights/biases.
            batch_size (int): Number of samples per batch.
            weight_reg (float): Multiplier for L2 penalty on weights.
            display_step (int): Cost is printed every display_steps during training.
            save_files (bool): If True, saves train/valid cost per epoch, weights/biases.
            file_name (str): File name used if 'save_files' is True.
            save_model (bool): If True, saves the trained model.
            model_name (str): File name for model if 'save_model' is True.
            model_path (str): Location of where the model/files are saved.
            random_state (int): Set the random seed.

        Returns:
            None
        """

        if(learn_rate == None):
            optimal=self.get_hyperparameters(X, Y, weight_reg, random_state)
            learn_rate=optimal['learn_rate']
            batch_size=optimal['batch_size']
            epochs=optimal['epochs']

        self.train(X, Y, epochs, learn_rate, batch_size, weight_reg,
              display_step, save_files, file_name, save_model,
              model_name, model_path, random_state)



    def load_model(self, model_name, model_path):
        """Reload tensorflow session for saved model. Called by djinn.load,

        Args:
            model_path (str, optional): Location of model if different than
                       location set during training.
            model_name (str, optional): Name of model if different than
                       name set during training.

        Returns:
            Object: djinn regressor model.

        """
        self.__sess = {}
        for p in range(0, self.__n_trees):
            tf.reset_default_graph()
            new_saver = \
            tf.train.import_meta_graph('%s%s_tree%s.ckpt.meta'%(model_path,model_name,p))
            self.__sess[p] = tf.Session()
            new_saver.restore(self.__sess[p], '%s%s_tree%s.ckpt'%(model_path,model_name,p))
            print("Model %s restored"%p)



    def close_model(self):
        """Closes tensorflow sessions launched with djinn.load.

        Args:
            None
        Returns:
            None
        """
        for p in range(0, self.__n_trees):
            self.__sess[p].close()



    def bayesian_predict(self, x_test, n_iters, random_state=None):
        """Bayesian distribution of predictions for a set of test data.

        Args:
            x_test (ndarray): Input parameters for test data.
            n_iters (int): Number of times to evaluate each neural network
                           per test point.
            random_state (int): Set the random seed.

        Returns:
            tuple (ndarray, ndarray, ndarray, dict):
                25th percentile of distribution of predictions for each test point.
                50th percentile of distribution of predictions for each test point.
                75th percentile of distribution of predictions for each test point.
                Dictionary containing inputs and predictions per tree, per
                  iteration, for each test point.

        """
        nonBayes = False
        if(n_iters == None):
            nonBayes = True
            n_iters = 1
        if(random_state): tf.set_random_seed(random_state)
        if(self.__sess == None): self.load_model(self.modelname, self.modelpath)
        if(x_test.ndim == 1): x_test = x_test.reshape(1,-1)
        samples = {}
        samples['inputs'] = x_test
        x_test = self.__xscale.transform(x_test)
        samples['predictions'] = {}
        for p in range(0, self.__n_trees):
            x = self.__sess[p].graph.get_tensor_by_name("input:0")
            keep_prob = self.__sess[p].graph.get_tensor_by_name("keep_prob:0")
            pred = self.__sess[p].graph.get_tensor_by_name("prediction:0")
            samples['predictions']['tree%s'%p] = \
                   [self.__yscale.inverse_transform(self.__sess[p].run(pred,\
                   feed_dict={x:x_test, keep_prob:self.__dropout_keep_prob}))
                   for i in range(n_iters)]

        nout = samples['predictions']['tree0'][0].shape[1]
        preds = np.array([samples['predictions'][t]
                for t in samples['predictions']]).reshape((n_iters*self.__n_trees, len(x_test), nout))

        middle = np.percentile(preds, 50, axis=0)
        lower = np.percentile(preds, 25, axis=0)
        upper = np.percentile(preds, 75, axis=0)
        if(nonBayes == True):
            return(np.mean(preds, axis=0))
        else:
            return(lower, middle, upper, samples)

    def gradient(self,x_test,input_idx=None,output_idx=0):
        print('calculating gradient...')

        grs = []
        for p in range(0, self.__n_trees):
            print('TREE {}:'.format(p))
            x = self.__sess[p].graph.get_tensor_by_name("input:0")
            pred = self.__sess[p].graph.get_tensor_by_name("prediction:0")
            unscale = self.__yscale.scale_[output_idx]
            if input_idx==None:
                gr = tf.gradients(pred[:,output_idx],x,stop_gradients=[x])
            else:
                gr = tf.gradients(pred[:,output_idx],x,stop_gradients=x)

            keep_prob = self.__sess[p].graph.get_tensor_by_name("keep_prob:0")
            x_test1 = self.__xscale.transform(x_test)
            val = self.__sess[p].run(gr[0],\
                feed_dict={x:x_test1, keep_prob:self.__dropout_keep_prob})
            if input_idx==None:
                val = val*self.__xscale.scale_[0]/unscale
            else:
                print(val)
                val = val[:,input_idx]*self.__xscale.scale_[input_idx]/unscale
            grs.append(val)
        return grs


    def gradient_v2(self,x_test,input_idx=None,output_idx=0):
        ''' BROKEN VERSION, EXPERIMENTING WITH GRADIENT TAPE '''
        print('calculating gradient...')

        grs = []
        for p in range(0, self.__n_trees):
            print('TREE {}:'.format(p))
            x = self.__sess[p].graph.get_tensor_by_name("input:0")
            pred = self.__sess[p].graph.get_tensor_by_name("prediction:0")
            unscale = self.__yscale.scale_[output_idx]
            # if input_idx==None:
            #     gr = tf.gradients(pred[:,output_idx],x)
            # else:
            #     gr = tf.gradients(pred[:,output_idx],x[:,input_idx])

            keep_prob = self.__sess[p].graph.get_tensor_by_name("keep_prob:0")
            x_test1 = self.__xscale.transform(x_test)
            with tf.GradientTape() as tape:
                tape.watch(x)
                f = self.__yscale.inverse_transform(self.__sess[p].run(pred,\
                feed_dict={x:x_test, keep_prob:self.__dropout_keep_prob}))
            grad_f = tape.gradient(f, x)
            val = grad_f[0, 0]
            # if input_idx==None:
            #     val = val*self.__xscale.scale_[0]/unscale
            # else:
            #     val = val*self.__xscale.scale_[input_idx]/unscale
            grs.append(val)
        return grs

    def hessian(self,x_test):
        print('calculating hessian...')

        hss =[]
        for p in range(0, self.__n_trees):
            print('TREE {}:'.format(p))
            x = self.__sess[p].graph.get_tensor_by_name("input:0")
            pred = self.__sess[p].graph.get_tensor_by_name("prediction:0")
            hs = tf.hessians(pred,x)

            keep_prob = self.__sess[p].graph.get_tensor_by_name("keep_prob:0")
            val = self.__sess[p].run(hs[0],\
                feed_dict={x:x_test, keep_prob:self.__dropout_keep_prob})
            hss.append(val)
        return hss


    def predict(self, x_test, random_state=None):
        """Predict target values for a set of test data.

        Args:
            x_test (ndarray): Input parameters for test data.
            random_state (int): Set the random seed.

        Returns:
            ndarray: Mean target value prediction for each test point.
        """
        return self.bayesian_predict(x_test, None, random_state)


    def collect_tree_predictions(self, predictions):
        """Gather distributions of predictions for each test point.

        Args:
            predictions (dict): The 'predictions' key from the dictionary
                        returned by bayesian_predict.

        Returns:
            ndarray: Re-shaped predictions (niters*ntrees, # test points, output dim)
        """
        nout = predictions['tree0'][0].shape[1]
        n_iters = len(predictions['tree0'])
        xlength = predictions['tree0'][0].shape[0]
        preds = np.array([predictions[t]
                for t in predictions]).reshape((n_iters*self.__n_trees, xlength, nout))
        return(preds)


    def continue_training(self, X, Y, training_epochs, learn_rate, batch_size, display_step=1, random_state=None):
        """Continue training an exisitng model. Must load_model first.

        Model is resaved in current location.

        Args:
            X (ndarray): Input parameters for training data.
            Y (ndarray): Target parameters for training data.
            epochs (int): Number of training epochs.
            learn_rate (float): Learning rate for optimizaiton of weights/biases.
            batch_size (int): Number of samples per batch.
            display_step (int): Cost is printed every display_steps during training.
            random_state (int): Set the random seed.

        Returns:
            None
        """
        ntrees=self.__n_trees
        nhl=self.__tree_max_depth-1
        dropout_keep_prob=self.__dropout_keep_prob
        tf_continue_training(self.__regression, self.__xscale, self.__yscale,
                            X, Y, ntrees, learn_rate, training_epochs, batch_size,
                            self.__dropout_keep_prob, nhl, display_step,
                            self.modelname, self.modelpath, random_state)




class DJINN_Classifier(DJINN_Regressor):
    """DJINN classification model.

    Args:
        n_trees (int): Number of trees in random forest
                (equal to the number of neural networks).
        max_tree_depth (int): Maximum depth of decision tree.
                       Neural network will have max_tree_depth-1
                       hidden layers.
        dropout_keep_prob (float): Probability of keeping a neuron
                          in dropout layers.
    """
    def __init__(self, n_trees=1, max_tree_depth=4, dropout_keep_prob=1.0):
        DJINN_Regressor.__init__(self, n_trees, max_tree_depth, dropout_keep_prob)
        self._DJINN_Regressor__regression = False

    def bayesian_predict(self, x_test, n_iters, random_state=None):
        """Bayesian distribution of predictions for a set of test data.

        Args:
            x_test (ndarray): Input parameters for test data.
            n_iters (int): Number of times to evaluate each neural network
                           per test point.
            random_state (int): Set the random seed.

        Returns:
            tuple (ndarray, ndarray, ndarray, dict):
                25th percentile of distribution of predictions for each test point.
                50th percentile of distribution of predictions for each test point.
                75th percentile of distribution of predictions for each test point.
                Dictionary containing inputs and predictions per tree, per
                  iteration, for each test point.

        """
        nonBayes = False
        if(n_iters == None):
            nonBayes = True
            n_iters = 1
        if(random_state): tf.set_random_seed(random_state)
        if(self._DJINN_Regressor__sess == None): self.load_model(self.modelname, self.modelpath)
        if(x_test.ndim == 1): x_test = x_test.reshape(1,-1)
        samples = {}
        samples['inputs'] = x_test
        x_test = self._DJINN_Regressor__xscale.transform(x_test)
        samples['predictions'] = {}
        for p in range(0, self._DJINN_Regressor__n_trees):
            x = self._DJINN_Regressor__sess[p].graph.get_tensor_by_name("input:0")
            keep_prob = self._DJINN_Regressor__sess[p].graph.get_tensor_by_name("keep_prob:0")
            pred = self._DJINN_Regressor__sess[p].graph.get_tensor_by_name("prediction:0")
            samples['predictions']['tree%s'%p] = \
                  [self._DJINN_Regressor__sess[p].run(pred,\
                  feed_dict={x:x_test, keep_prob:self._DJINN_Regressor__dropout_keep_prob}) for i in range(n_iters)]

        nout = samples['predictions']['tree0'][0].shape[1]
        preds = np.array([samples['predictions'][t]
                for t in samples['predictions']]).reshape((n_iters*self._DJINN_Regressor__n_trees, len(x_test), nout))
        print(preds.shape)
        middle = np.argmax(np.percentile(preds,50,axis=0),1)
        lower = np.argmax(np.percentile(preds,25,axis=0), 1)
        upper = np.argmax(np.percentile(preds,75,axis=0), 1)

        if(nonBayes == True):
            return(middle)
        else:
            return(lower, middle, upper, samples)
