import theano
import theano.tensor as T
import numpy as np
import random
import matplotlib.pyplot as plt

import pycof as pc

from .optimizers import Adam, AdaGrad, AdaMax
from .optimizers import SGD, MomentumSGD
from .optimizers import RMSprop
from .initializations import init_params
from .activations import tanh, sigmoid, relu, softplus
from .losses import mean_squared_error, binary_cross_entropy

"""
DeepLearning website - Yoshua Bengio:
- http://deeplearning.net/tutorial/mlp.html
"""

## TODO: add batch
## TODO: comment out

T_activations = {None: None,
                'linear': None,
                'tanh': T.tanh,
                'sigmoid': T.nnet.sigmoid,
                'relu': T.nnet.relu,
                'softplus': T.nnet.softplus,
                }

activations = {None: None,
                'linear': None,
                'tanh': tanh,
                'sigmoid': sigmoid,
                'relu': relu,
                'softplus': softplus,
                }

optimizers = {'sgd': SGD,
                'stochastic_gradient_descent': SGD,
                'momentum': MomentumSGD,
                'momentumsgd': MomentumSGD,
                'adam': Adam,
                'adagrad': AdaGrad,
                'adamax': AdaMax,
                'rmsprop': RMSprop,
                }


class Layer(object):
    def __init__(self, n_in, n_out, W=None, b=None, activation=None, seed=None, init_weights='xavier', init_bias='zeros'):
        """Fully connected Neural Network layer
        
        :param n_in: Dimension of the input layer (number of rows for the weights matrix).
        :type n_in: int
        :param n_out: Dimension of the output layer (number of columns for the weights matrix).
        :type n_out: int
        :param W: Weights matrix, defaults to None (generated by :obj:`init_params` in `initializations <initializations.html#statinf.ml.initializations.init_params>`_).
        :type W: numpy.array, optional
        :param b: Vector of bias, defaults to None (generated by :obj:`init_params` in `initializations <initializations.html#statinf.ml.initializations.init_params>`_).
        :type b: numpy.array, optional
        :param activation: Activation function to be applied, defaults to None (i.e. linear, see `activations <activations.html>`_)
        :type activation: str, optional
        :param seed: Random seed, defaults to None.
        :type seed: int, optional
        :param init_weights: Distribution to be used when initializing the weights matrix, defaults to 'xavier' (see `initializations <initializations.html#statinf.ml.initializations.init_params>`_).
        :type init_weights: str, optional
        :param init_bias: Method to used for initializing the bias vector, defaults to 'zeros' (see `initializations <initializations.html#statinf.ml.initializations.init_params>`_).
        :type init_bias: str, optional

        :return: Layer to be stacked to the Neural Network
        :rtype: Layer
        """
      
        self.input = input
        self.activation = activation
        self.f = T_activations[activation]
        self.activate = activations[activation]

        if W is None:
            # If weights matrix is not provided (default), initialize from a distribution
            W_values = init_params(rows=n_in, cols=n_out, method=init_weights, seed=seed)
            # Multiply by 4 the weights for sigmoid activation. See Y. Bengio's website
            if activation == theano.tensor.nnet.sigmoid:
                W_values *= 4
        elif W.shape == (n_in, n_out):
            raise ValueError(f'Weights dimension does not match. Dimension for W should be {(n_in, n_out)} and got {W.shape}.')

        if b is None:
            # If bias matrix is not provided (default), initialize from a distribution
            b_values = init_params(rows=1, cols=n_out, method=init_bias, seed=seed)
        elif b.shape == (n_in, n_out):
            raise ValueError(f'Weights dimension does not match. Dimension for b should be {(1, n_out)} and got {b.shape}.')
        
        # Store parameters in the class as shared parameters with theano
        self.W = theano.shared(value=W_values, name='W', borrow=True)
        self.b = theano.shared(value=b_values, name='b', borrow=True)
        self.optimal_w = W_values
        self.optimal_b = b_values
        
        # Parameters of the model
        self.params = [self.W, self.b]

    def _feed_forward(self, input, tensor=True):
        if tensor:
            Xb_e = T.dot(input, self.W) + self.b
            self.output = (Xb_e if self.activation in [None, 'linear'] else self.f(Xb_e))
        else:
            # Used as feed forward for prediction (avoid using Theano)
            Xb_e = input.dot(self.W.get_value()) + self.b.get_value()
            self.output = (Xb_e if self.activation in [None, 'linear'] else self.activate(Xb_e))
        return self.output

class MLP:
    """Multi-Layer Perceptron
    
    :references: Rosenblatt, F. (1958). `The perceptron: a probabilistic model for information storage and organization in the brain <http://web.engr.oregonstate.edu/~huanlian/teaching/machine-learning/2017fall/extra/rosenblatt-1958.pdf>`_. Psychological review, 65(6), 386.
    
    """
    def __init__(self, loss='MSE', random=None):
        self.loss = loss
        self._layers = []
        self.params = []
        self.L1 = 0.
        self.L2 = 0.
        self._cost_value = []

    def _L1(self):
        """Compute L1 regularization
        
        :formula: .. math:: L_{1} = \sum_{i=1}^{L} |\\beta|

        :return: L1 penalty
        :rtype: float
        """
        for layer in self._layers:
            # L1 norm ; one regularization option is to enforce L1 norm to
            # be small
            self.L1 += abs(layer.W).sum()
        return self.L1
    
    def _L2(self):
        """Compute L2 regularization
        
        :formula: .. math:: L_{2} = \sum_{i=1}^{L} \\beta^{2}

        :return: L2 penalty
        :rtype: float
        """
        for layer in self._layers:
            # square of L2 norm ; one regularization option is to enforce
            # square of L2 norm to be small
            self.L2 += (layer.W ** 2).sum()
        return self.L2
    
    def _params(self):
        # the parameters of the model are the parameters of the two layer it is
        # made out of
        _params = []
        for layer in self._layers:
            _params += list(layer.params)
        
        self.params = list(_params)
    
    def add(self, layer):
        """Stack layer to Neural Network
        
        :param layer: Layer to be stacked
        :type layer: :func:`~Layer`
        """
        self._layers.append(layer)
    
    def _forward_prop(self, x, tensor=True):
        output = x
        for layer in self._layers:
            output = layer._feed_forward(output, tensor)
        return output

    def _cost(self, x, y):
        output = self._forward_prop(x)
        # Compute loss
        if self.loss.lower() in ['mse', 'mean_squared_error']:
            _cost_value = mean_squared_error(y_true=y, y_pred=output) # ((output - y) ** 2).sum()
        elif self.loss.lower() in ['binary_cross_entropy', 'bce']:
            _cost_value = binary_cross_entropy(y_true=y, y_pred=output)
        else:
            raise ValueError('Loss function is not valid.')
        self._cost_value += [_cost_value]
        return _cost_value
    

    def get_weights(self, layer='all', param='all'):
        """Fetches the parameters from the network.
        
        :param layer: Layer id from which to fetch the parameters, defaults to 'all'
        :type layer: int, optional
        :param param: What parameter we need to fetch (can either be 'weights', 'bias' or 'all'), defaults to 'all'
        :type param: str, optional
        
        :return: Weights and Bias used in the Neural Network.
        :rtype: dict
        """
        # Check what parameter the user was to get
        if param.lower() in ['weights', 'w', 'weight']:
            par = 0
        elif param.lower() in ['biases', 'b', 'bias']:
            par = 1
        elif param.lower() == 'all':
            par = 2
        else:
            raise ValueError('Value for "param" is not value. Please chose between "weights", "bias" or "all".')
        
        if layer == 'all':
            # If user wants to see all layers, we create a dictionnary
            weights = {}
            for i in range(len(self._layers)):
                if par == 0:
                    weights.update({f'weights {i}': self._layers[i].W.get_value()})
                elif par == 1:
                    weights.update({f'bias {i}': self._layers[i].b.get_value()})
                else:
                    weights.update({f'weights {i}': self._layers[i].W.get_value(), f'bias {i}': self._layers[i].b.get_value()})
        elif layer in range(len(self._layers)):
            # If user wants only 1 specific layer,
            if par == 2:
                # we return a dict for all params
                weights = {'weights': self._layers[layer].Z.get_value(), 'bias': self._layers[layer].b.get_value()}
            else:
                # or an array for 1 single param
                weights = self._layers[layer].Z.get_value()
        else:
            raise ValueError(f'Layer is incorrect. Please chose either "all" or layer <= {len(self._layers) - 1}. Got layer = {layer}')
        
        return weights

    def train(self, data, X, Y='Y', epochs=100, optimizer='SGD', batch_size=1, training_size=0.8, test_set=None, learning_rate=0.01, L1_reg=0., L2_reg=0., early_stop=True, patience=10, improvement_threshold=0.995, restore_weights=True, verbose=True, verbose_all=False, plot=False):
        """Train the Neural Network.
        
        :param data: Layer id from which to fetch the parameters.
        :type data: pandas.DataFrame
        :param X: List of X input variables.
        :type X: list
        :param Y: Variable to predict, defaults to 'Y'.
        :type Y: str, optional
        :param epochs: Number of epochs to train the network, defaults to 100.
        :type epochs: int, optional
        :param optimizer: Algortihm to use to minimize the loss function, defaults to 'SGD' (see `optimizers <optimizers.html>`_).
        :type optimizer: str, optional
        :param batch_size: Size of each batch to be trained, defaults to 1.
        :type batch_size: int, optional
        :param training_size: Ratio of the data to be used for training set (:math:`\in (0, 1)`) the remainder is used for test set, defaults to 0.8.
        :type training_size: float, optional
        :param test_set: Data frame to use as test set (overrides **training_size** if provided), defaults to None.
        :type test_set: pandas.DataFrame, optional
        :param learning_rate: Learning rate (step size) for gradient descent, defaults to 0.01.
        :type learning_rate: float, optional
        :param L1_reg: Coefficient :math:`\\lambda_{1} \in (0,1)` for the L1 penalty, defaults to 0.
        :type L1_reg: float, optional
        :param L2_reg: Coefficient :math:`\\lambda_{2} \in (0,1)` for the L2 penalty, defaults to 0.
        :type L2_reg: float, optional
        :param early_stop: Apply early-stopping method to avoid over-fitting. Stop the training sequence after the patience (see below) with no improvement, defaults to True.
        :type early_stop: bool, optional
        :param patience:  Number of iterations to wait after a minimum has been found. If no improvement, training sequence is stopped, defaults to 10.
        :type patience: int, optional
        :param improvement_threshold: Change level of the loss from which we consider it is an improvement, defaults to 0.995.
        :type improvement_threshold: float, optional
        :param restore_weights: If training is early-stopped, restore the weights of the best iteration, defaults to True.
        :type restore_weights: bool, optional
        :param verbose: Verbosity level (True or 1 for minimal display, 2 for details), defaults to True.
        :type verbose: int, optional
        :param plot: Plot the evolution of the training and test loss for each epoch, defaults to False.
        :type plot: bool, optional

        :references: * Bergstra, J., Bastien, F., Breuleux, O., Lamblin, P., Pascanu, R., Delalleau, O., ... & Bengio, Y. (2011). `Theano: Deep learning on gpus with python <http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.678.1889&rep=rep1&type=pdf>`_. In NIPS 2011, BigLearning Workshop, Granada, Spain (Vol. 3, pp. 1-48). Microtome Publishing. 
            * Friedman, J., Hastie, T., & Tibshirani, R. (2001). `The elements of statistical learning <https://web.stanford.edu/~hastie/Papers/ESLII.pdf>`_ (Vol. 1, No. 10). New York: Springer series in statistics.
            * Goodfellow, I., Bengio, Y., & Courville, A. (2016). `Deep learning <https://www.deeplearningbook.org/>`_. MIT press. 
        """
        # Parse arguments to pass to self
        self.learning_rate = learning_rate
        self.L1_reg = L1_reg
        self.L2_reg = L2_reg
        self.X_col = X
        self.Y_col = Y
        self.optimizer = str(optimizer).lower()

        verbose = True if verbose in [True, 1, 2] else False
        verbose_all = True if verbose in [2] else False

        # Get indexes
        self.data_size = data.shape[0]

        if test_set is None:
            data_idx = list(data.index)
            train_index = random.sample(data_idx, int(self.data_size * training_size))
            test_index = [i for i in data_idx if i not in train_index]
            # Train set
            train_set_x = data.loc[train_index, self.X_col]
            train_set_y = data.loc[train_index, self.Y_col]
            # Test set
            test_set_x = data.loc[test_index, self.X_col]
            test_set_y = data.loc[test_index, self.Y_col]
        else:
            # Train set
            train_set_x = data[self.X_col]
            train_set_y = data[self.Y_col]
            # Test set
            test_set_x = test_set[self.X_col]
            test_set_y = test_set[self.Y_col]
        
        # Set data to theano
        # Train set
        train_set_x = theano.shared(np.array(train_set_x, dtype=theano.config.floatX))
        train_set_y = theano.shared(np.array(train_set_y, dtype=theano.config.floatX))
        # Test set
        test_set_x = theano.shared(np.array(test_set_x, dtype=theano.config.floatX))
        test_set_y = theano.shared(np.array(test_set_y, dtype=theano.config.floatX))
        self.test_set = test_set_x

        index = T.lscalar()  # index to a [mini]batch
        x = T.vector('x')
        y = T.scalar('y')

        # Get initial parameters
        self._params()

        # Set cost
        cost = self._cost(x, y) + (self.L1_reg * self._L1()) + (self.L2_reg * self._L2())
        
        ## Define update with optimizer
        try:
            self.updates = optimizers[self.optimizer](params=self.params, learning_rate=self.learning_rate).updates(cost)
        except BaseException:
            posible_opt = "', '".join([i for i in optimizers.keys()])
            raise ValueError(f"Optimizer not value. Please make sure you choose in '{posible_opt}'. Got '{self.optimizer}'.")
 
        # Define Theano function for training step
        train_model = theano.function(
            inputs=[index],
            outputs=cost,
            updates=self.updates,
            givens={x: train_set_x[index], y: train_set_y[index]}
        )
        
        # Define Theano function for validation step
        validate_model = theano.function(
            inputs=[index],
            outputs=cost,
            givens={x: test_set_x[index], y: test_set_y[index]}
        )
        # Define length of test
        n_train_set = train_set_x.get_value(borrow=True).shape[0]
        n_test_set = test_set_x.get_value(borrow=True).shape[0]
        
        # Initialize values
        self.train_losses = []       # Training loss vector to plot
        self.test_losses = []        # Test loss vector to plot
        self.best_loss = np.inf # Initial best loss value
        self.best_iter = 0
        verif_i = epochs        # First, verify epoch at then end (init only)
        i = 0

        # Start training
        while i <= epochs:
            i += 1
            # Measure accuracy on train set
            epoch_loss = np.mean([train_model(i) for i in range(n_train_set)])
            # Evaluate on test set
            epoch_loss_test = np.mean([validate_model(i).tolist() for i in range(n_test_set)])
            # Save to historical performance
            self.train_losses += [epoch_loss]
            self.test_losses += [epoch_loss_test]
            # If we got the best score until now, and patience is not reached
            if (epoch_loss_test < self.best_loss * improvement_threshold) & (i <= verif_i) & early_stop:
                pc.verbose_display(f'New min found for iter {i}', verbose=verbose_all)
                verif_i = i + patience
                self.best_loss = epoch_loss_test
                self.best_iter = i
                # self.optimal_w = self.w.get_value()
                # self.optimal_b = self.b.get_value()
                for layer in self._layers:
                    layer.optimal_w = layer.W.get_value()
                    layer.optimal_b = layer.b.get_value()
            # if no improvement
            elif (i < verif_i):
                # pc.verbose_display(f'No improvement at iter {i}', verbose=verbose_all)
                pass
            else:
                pc.verbose_display(f'Stop training. Minimum found after {self.best_iter} iterations', verbose=verbose)
                if restore_weights:
                    for layer in self._layers:
                        layer.W = theano.shared(value=layer.optimal_w, name='W', borrow=True)
                        layer.b = theano.shared(value=layer.optimal_b, name='b', borrow=True)
                        layer.params = [layer.W, layer.b]
                break

        if plot:
            # Plot the training and test loss
            plt.title('MLP train vs. test error through epochs', loc='center')
            plt.plot(self.train_losses, label='Training loss')
            plt.plot(self.test_losses, label='Test loss')
            plt.legend()
            plt.show()
    

    def predict(self, new_data, binary=False, threshold=0.5):
        """
        """
        # Transform input data to be transformed into numpy array format
        x = new_data[self.X_col].values
        # Apply feedforward step to the data to get the prediction (apply weights and activations)
        output = self._forward_prop(x, tensor=False)
        if binary:
            return [1. if y[0] > threshold else 0. for y in output.reshape((new_data.shape[0], 1))]
        else:
            return [y[0] for y in output.reshape((new_data.shape[0], 1))]
