# -*- coding: utf-8 -*-
"""
Created on Fri Aug 24 12:47:19 2018

@author: nkdi
"""
import numpy as np
# from w2l import activation_functions
# from surrogates import math_utils
from hipersim.surrogates import math_utils


class ann:
    def __init__(self, **kwargs):
        # print(kwargs)
        if 'layersizes' in kwargs:
            self.layersizes = kwargs['layersizes']
        else:
            self.layersizes = [1, 10, 1]

        self.nlayers = len(self.layersizes) - 1  # Input layer is not counted

        if 'neurontypes' in kwargs:
            self.neurontypes = kwargs['neurontypes']
        else:
            self.neurontypes = ['Linear'] * self.nlayers

        if 'activation' in kwargs:
            self.activation = kwargs['activation']
        else:
            self.activation = ['Relu'] * (self.nlayers - 1) + ['Linear']

        if 'testratios' in kwargs:
            self.testratios = kwargs['testratios']
        else:
            self.testratios = [0.7, 0.2, 0.1]

        if 'output_style' in kwargs:
            self.output_style = kwargs['output_style']
        else:
            self.output_style = 'None'

        if 'params' in kwargs:
            self.params = kwargs['params']
        else:
            self.params = {'weightinit': [0.01] * self.nlayers,
                           'regularization': 0.1,
                           'learningrate': 0.1,
                           'costfunction': 'Lsq',
                           'minibatchsize': 0,
                           'maxiter': 1500,
                           'RMSprop': 0.999,
                           'Momentum': 0.9,
                           'Xscaling': True,
                           'Yscaling': True}
        # Parse parameters that may have not been given in the input
        if 'weightinit' not in self.params:
            self.params['weightinit'] = [0.01] * self.nlayers
        if 'regularization' not in self.params:
            self.params['regularization'] = 0.1
        if 'learningrate' not in self.params:
            self.params['learningrate'] = 0.1
        if 'costfunction' not in self.params:
            self.params['costfunction'] = 'Lsq'
        if 'minibatchsize' not in self.params:
            self.params['minibatchsize'] = 0
        if 'maxiter' not in self.params:
            self.params['maxiter'] = 1500  # Same as nepochs
        if 'RMSprop' not in self.params:
            self.params['RMSprop'] = 0.999
        if 'Momentum' not in self.params:
            self.params['Momentum'] = 0.9
        if 'nepochs' not in self.params:
            self.params['nepochs'] = self.params['maxiter']
        if 'Xscaling' not in self.params:
            self.params['Xscaling'] = True
        if 'Yscaling' not in self.params:
            self.params['Yscaling'] = True

    def initialize_weights(self):
        W = []
        b = []
        # np.random.seed(1) # Used only in case we need to maintan consisteny with Matlab
        for i in range(self.nlayers):
            if self.neurontypes[i] == 'Linear':
                # Wi =
                # np.random.randn(self.layersizes[i],self.layersizes[i+1])*self.params['weightinit'][i]
                # # Regular way of initializing weights
                Wi = math_utils.NormalDist(2,
                                           np.random.rand(self.layersizes[i],
                                                          self.layersizes[i + 1]),
                                           0,
                                           1) * self.params['weightinit'][i]  # For consistency with Matlab
                bi = np.zeros([1, self.layersizes[i + 1]])
            elif self.neurontypes[i] == 'LinearPretrained':
                Wi = self.weights['W'][i]
                bi = self.weights['bias'][i]
            elif self.neurontypes[i] == 'LinearNonTrainable':
                Wi = self.weights['W'][i]
                bi = self.weights['bias'][i]
            W.append(Wi)
            b.append(bi)

        if hasattr(self, 'weights'):
            self.weights['W'] = W
            self.weights['bias'] = b
        else:
            self.weights = {'W': W, 'bias': b}

        return W, b

    def forward_propagation(self, X):
        A = []
        Z = []
        A.append(X)
        Ai = X
        for i in range(self.nlayers):
            if self.neurontypes[i] == 'Linear':
                Zi = np.dot(Ai, self.weights['W'][i]) + self.weights['bias'][i]
            elif self.neurontypes[i] == 'LinearPretrained':
                Zi = np.dot(Ai, self.weights['W'][i]) + self.weights['bias'][i]
            elif self.neurontypes[i] == 'LinearNonTrainable':
                Zi = np.dot(Ai, self.weights['W'][i]) + self.weights['bias'][i]
            else:
                raise ValueError('Unknown or unspecified neuron type')

            if self.activation[i] == 'Relu':
                Ai = math_utils.relu(Zi)
            elif self.activation[i] == 'Linear':
                Ai = Zi
            elif self.activation[i] == 'Tanh':
                Ai = math_utils.hyperbolic_tangent(Zi)
            elif self.activation[i] == 'Softplus':
                Ai = math_utils.softplus(Zi)
            else:
                raise ValueError('Unknown or unspecified activation function')

            A.append(Ai)
            Z.append(Zi)

        Ypredict = Ai
        ANNcache = (A, Z)
        return Ypredict, ANNcache

    def backward_propagation(self, X, Y, Ypredict, ANNcache):
        if self.params['costfunction'] == 'Log':
            dJdAprev = -Y / Ypredict + (1 - Y) / (1 - Ypredict)
        elif self.params['costfunction'] == 'Lsq':
            dJdAprev = (Ypredict - Y)
        else:
            raise ValueError('Unknown cost function type')

        dJdA = [None] * self.nlayers
        dAdZ = [None] * self.nlayers
        dJdW = [None] * self.nlayers
        dJdb = [None] * self.nlayers
        dJdZ = [None] * self.nlayers
        A, Z = ANNcache

        for i in range(self.nlayers - 1, -1, -1):
            dJdA[i] = dJdAprev
            if self.activation[i] == 'Linear':
                dAdZ[i] = np.ones(Z[i].shape)
            elif self.activation[i] == 'Relu':
                dAdZ[i] = np.ones(Z[i].shape) * (Z[i] > 0)
            elif self.activation[i] == 'Softplus':
                dAdZ[i] = math_utils.logistic(Z[i])
            elif self.activation[i] == 'Tanh':
                dAdZ[i] = 1 - math_utils.hyperbolic_tangent(Z[i])**2
            elif self.activation[i] == 'Logistic':
                dAdZ[i] = math_utils.logistic(Z[i]) * (1 - math_utils.logistic(Z[i]))
            else:
                raise ValueError('Unknown activation function')

            dJdZ[i] = dAdZ[i] * dJdA[i]
            mlength = X.shape[0]
            if self.neurontypes[i] == 'Linear':
                dJdW[i] = (1 / mlength) * np.dot(A[i].T, dJdZ[i])
                dJdb[i] = (1 / mlength) * np.sum(dJdZ[i], 0, keepdims=True)
            elif self.neurontypes[i] == 'LinearPretrained':
                dJdW[i] = (1 / mlength) * np.dot(A[i].T, dJdZ[i])
                dJdb[i] = (1 / mlength) * np.sum(dJdZ[i], 0, keepdims=True)
            elif self.neurontypes[i] == 'LinearNonTrainable':
                dJdW[i] = np.zeros(Z[i].shape)
                dJdb[i] = np.zeros((1, Z[i].shape[1]))
            else:
                raise ValueError('Unknown neuron type')

            dJdAprev = np.dot(dJdZ[i], self.weights['W'][i].T)

        return dJdW, dJdb

    def backward_propagation_dy(self, X):
        if self.params['Xscaling']:
            X = (X - self.Xbias) / self.Xscale  # Lots of broadcasting taking place here

        Y, ANNcache = self.forward_propagation(X)
        if self.params['Yscaling']:
            Y = Y * self.Yscale + self.Ybias

        A, Z = ANNcache
        noutputs = Y.shape[1]

        dY = [None] * noutputs

        for iY in range(noutputs):

            for i in range(self.nlayers - 1, -1, -1):

                if i == self.nlayers - 1:
                    Zi = Z[i][:, iY]
                    Zi = Zi.reshape(Zi.shape[0], 1)
                    Wi = self.weights['W'][i][:, iY]
                    Wi = Wi.reshape(Wi.shape[0], 1)
                else:
                    Zi = Z[i]
                    Wi = self.weights['W'][i]

                if self.activation[i] == 'Linear':
                    dAdZ = np.ones(Zi.shape)
                elif self.activation[i] == 'Relu':
                    dAdZ = np.ones(Zi.shape) * (Zi > 0)
                elif self.activation[i] == 'Softplus':
                    dAdZ = math_utils.logistic(Zi)
                elif self.activation[i] == 'Tanh':
                    dAdZ = 1 - math_utils.hyperbolic_tangent(Zi)**2
                elif self.activation[i] == 'Logistic':
                    dAdZ = math_utils.logistic(Zi) * (1 - math_utils.logistic(Zi))
                else:
                    raise ValueError('Unknown activation function')

                if i == self.nlayers - 1:
                    dYi = np.dot(dAdZ, Wi.T)
                else:
                    dYi = np.dot((dAdZ * dYi), Wi.T)

            if self.params['Yscaling']:
                dYi = dYi * self.Yscale[0, iY]
            if self.params['Xscaling']:
                dYi = dYi / self.Xscale

            dY[iY] = dYi

        return Y, dY

    def check_gradients(self, X):
        if not hasattr(self, 'weights'):
            self.initialize_weights()

    def check_gradients_dxdy(self, X):
        Yout, dYout = self.backward_propagation_dy(X)

    def compute_cost(self, X, Y):
        # Computes the cost using log-functions
        mtrain = X.shape[0]
        noutputs = Y.shape[1]
        Ypredict, ANNcache = self.forward_propagation(X)
        if self.params['costfunction'] == 'Log':
            Loss = - (np.multiply(Y, np.log(Ypredict)) + np.multiply((1 - Y), np.log(1 - Ypredict)))
        elif self.params['costfunction'] == 'Lsq':
            Loss = 0.5 * np.multiply((Ypredict - Y), (Ypredict - Y))
        else:
            print('Unknown loss function type!')

        if self.params["regularization"] == 0:
            J = (1 / mtrain) * np.sum(Loss) * (1 / noutputs)
        else:
            v = cell_to_vector(self.weights['W'], self.weights['bias'])
            J = (1 / mtrain) * np.sum(Loss) * (1 / noutputs) + \
                (self.params['regularization'] / (2 * mtrain)) * (np.dot(v.T, v))

        return J, Ypredict, ANNcache

    def train(self, X, Y):
        """
        Uses input-output data pair (X,Y) to train the weights of a neural-net object
        """
        # Parse hyperparameters
        beta2 = self.params['RMSprop']
        if beta2 == 0:
            epsroundoff = 0
        else:
            epsroundoff = 1e-8
        beta1 = self.params['Momentum']
        minibatchsize = self.params['minibatchsize']
        nepochs = self.params['nepochs']

        self.Xbias = np.mean(X, axis=0, keepdims=True)  # Keepdims = True needed to allow broadcasting
        self.Ybias = np.mean(Y, axis=0, keepdims=True)
        self.Xscale = np.std(X, axis=0, ddof=1, keepdims=True)
        self.Yscale = np.std(Y, axis=0, ddof=1, keepdims=True)

        if self.params['Xscaling']:
            X = (X - self.Xbias) / self.Xscale
        if self.params['Yscaling']:
            Y = (Y - self.Ybias) / self.Yscale

        # Split data into training, validation and test sets
        Ntotal = X.shape[0]
        Ntrain = int(np.floor(Ntotal * self.testratios[0]))
        Ndev = int(np.floor(Ntotal * self.testratios[1]))
        Xtrain = X[:Ntrain, :]
        Ytrain = Y[:Ntrain, :]
        Xdev = X[Ntrain:Ntrain + Ndev, :]
        Ydev = Y[Ntrain:Ntrain + Ndev, :]
        Xtest = X[Ntrain + Ndev:, :]
        Ytest = Y[Ntrain + Ndev:, :]

        # Determine iteration scheme
        if minibatchsize == 0:  # Minibatch functionality off
            minibatchsize = Ntrain  # Each iteration is on the entire train set
            Nbatches = 1
            Niter = self.params['nepochs']
        else:
            Nbatches = int(np.ceil(Ntrain / minibatchsize))
            Niter = nepochs * Nbatches

        # Initialize variables and weights
        noutputs = Y.shape[1]
        Jhist = np.zeros([Niter, 1])
        Jgradnorm = np.zeros([Niter, 1])
        Jdevhist = np.zeros([Niter, 1])
        Rsqhist = np.zeros([Niter, noutputs])
        Rsqdevhist = np.zeros([Niter, noutputs])

        self.initialize_weights()
        ntotal = 0
        for i in range(self.nlayers):
            ntotal += self.weights['W'][i].shape[0] * self.weights['W'][i].shape[1] + \
                self.weights['bias'][i].shape[0] * self.weights['bias'][i].shape[1]

        # Initialize interactive plot/inline-output functionalities
        if self.output_style == 'Verbose':
            print('Verbose output')
        elif self.output_style == 'Plot':
            import tkinter
            import matplotlib
            import matplotlib.pyplot as plt
            plt.ion()
            matplotlib.use('TkAgg')
            fig, axs = plt.subplots(1, 2, figsize=(8, 4))
            plt00 = axs[0].plot(0, 0, '-k')[0]
            plt01 = axs[0].plot(0, 0, '-r')[0]
            axs[0].set_xlabel('Number of iterations')
            axs[0].set_ylabel('Cost function value J')
            plt10 = axs[1].plot(0, 0, '-k')[0]
            plt11 = axs[1].plot(0, 0, '-r')[0]
            plt12 = axs[1].plot(0, 0, '--k')[0]
            plt13 = axs[1].plot(0, 0, '--r')[0]

            axs[1].set_xlabel('Number of iterations')
            axs[1].set_ylabel('R-squared')
            axs[1].set_ylim([0, 1])
            fig.tight_layout()
            fig.show()
            fig.canvas.draw()

        elif self.output_style == 'None':
            print('Training neural net model:')

        # Begin iterations
        go = 1
        i = 0
        ibatch = 0
        iepoch = 0
        Vd = np.zeros([ntotal, 1])
        Sd = np.zeros([ntotal, 1])
        while go == 1:
            # Look at the dev(validation) set and compute cost function and R-square
            Jdev, Ypredictdev, _ = self.compute_cost(Xdev, Ydev)
            Jdevhist[i] = Jdev
            Rsqdevhist[i, :] = math_utils.pairwise_correlation(Ypredictdev, Ydev)**2

            # Pick a single minibatch of training data samples
            if ibatch == Nbatches:
                Xtraini = Xtrain[ibatch * minibatchsize:, :]
                Ytraini = Ytrain[ibatch * minibatchsize:, :]
            else:
                Xtraini = Xtrain[ibatch * minibatchsize:(ibatch + 1) * minibatchsize, :]
                Ytraini = Ytrain[ibatch * minibatchsize:(ibatch + 1) * minibatchsize, :]

            # Compute cost function, cost function gradients, and R-square on one minibatch
            J, Ypredict, ANNcache = self.compute_cost(Xtraini, Ytraini)  # Cost function
            [dJdW, dJdb] = self.backward_propagation(Xtraini, Ytraini, Ypredict, ANNcache)  # Cost function gradients
            dvecti = cell_to_vector(dJdW, dJdb)
            Jhist[i] = J
            # Gradient norm for monitoring purposes
            Jgradnorm[i] = (1 / minibatchsize) * np.sqrt(np.dot(dvecti.T, dvecti))
            Rsqhist[i, :] = math_utils.pairwise_correlation(Ypredict, Ytraini)**2  # R-square

            # Update weights based on gradient descent
            if (beta1 != 0) & (beta2 != 0):  # Adam algorithm
                Vd = beta1 * Vd + (1 - beta1) * dvecti
                Sd = beta2 * Sd + (1 - beta2) * (dvecti**2)
                Vcorr = Vd / (1 - beta1**(i + 1))
                Scorr = Sd / (1 - beta2**(i + 1))
                vupd = cell_to_vector(self.weights['W'], self.weights['bias']) - self.params['learningrate'] * (Vcorr / (np.sqrt(
                    Scorr) + epsroundoff) + cell_to_vector(self.weights['W'], self.weights['bias']) * self.params['regularization'] / Xtraini.shape[0])
            elif beta1 > 0 & beta2 == 0:  # Only momentum, no RMSprop, optional regularization
                Vd = beta1 * Vd + (1 - beta1) * dvecti
                Vcorr = Vd / (1 - beta1**(i + 1))
                vupd = cell_to_vector(self.weights['W'], self.weights['bias']) - self.params['learningrate'] * (
                    Vcorr + cell_to_vector(self.weights['W'], self.weights['bias']) * self.params['regularization'] / Xtraini.shape[0])
            else:
                # Simple gradient descent with optional regularization
                vupd = cell_to_vector(self.weights['W'],
                                      self.weights['bias']) - self.params['learningrate'] * (cell_to_vector(dJdW,
                                                                                                            dJdb) + cell_to_vector(self.weights['W'],
                                                                                                                                   self.weights['bias']) * self.params['regularization'] / Xtraini.shape[0])

            self.weights['W'], self.weights['bias'] = vector_to_cell(vupd, self.weights['W'], self.weights['bias'])

            # Print/plot iteration outputs
            if self.output_style == 'Verbose':
                print('Epoch #' + str(iepoch) + ', batch #' + str(ibatch) + ', cost function value = ' + str(J[0]))
            elif (self.output_style == 'Plot') & (i >= 1):

                # Update plot with results from current iteration
                xplot = np.arange(i)
                plt00.set_data(xplot, Jhist[:i])
                plt01.set_data(xplot, Jdevhist[:i])
                axs[0].relim()
                axs[0].autoscale_view(True, True, True)  # Autoscale
                plt10.set_data(xplot, np.min(Rsqhist[:i, :], 1))
                plt11.set_data(xplot, np.min(Rsqdevhist[:i, :], 1))
                plt12.set_data(xplot, np.max(Rsqhist[:i, :], 1))
                plt13.set_data(xplot, np.max(Rsqdevhist[:i, :], 1))

                axs[1].relim()
                axs[1].autoscale_view(True, True, True)  # Autoscale

                plt.pause(0.001)
                fig.canvas.draw()
                plt.show()

            i += 1
            ibatch += 1
            if i >= Niter:
                go = 0
            if ibatch >= Nbatches:
                if iepoch >= nepochs:
                    go = 0
                else:
                    ibatch = 0
                    iepoch += 1

        Jout, Yout, _ = self.compute_cost(Xtest, Ytest)
        Rsqout = math_utils.pairwise_correlation(Yout, Ytest)**2

        if self.params['Yscaling']:
            Testbias = (np.mean(Yout, 0) - np.mean(Ytest, 0)) * self.Yscale
        else:
            Testbias = np.mean(Yout, 0) - np.mean(Ytest, 0)

        self.metrics = {'RsqTest': Rsqout, 'RsqDev': Rsqdevhist[i -
                                                                1, :], 'RsqTrain': Rsqhist[i - 1, :], 'Testbias': Testbias}

        # Start from previous trained state in case a new training command is issued
        self.neurontypes = ['LinearPretrained' if ni == 'Linear' else ni for ni in self.neurontypes]

        Outdata = {'Jhist': Jhist, 'Yout': Yout, 'Jout': Jout}

        return Outdata

    def predict(self, X):
        """
        Run a model prediction (forward propagation + scaling)
        """
        if self.params['Xscaling']:
            X = (X - self.Xbias) / self.Xscale  # Lots of broadcasting taking place here

        Y, _ = self.forward_propagation(X)
        if self.params['Yscaling']:
            Y = Y * self.Yscale + self.Ybias  # Lots of broadcasting also taking place here

        return Y


def cell_to_vector(W, b):
    """
    Reshapes an array from dictionary-type structure to vector
    """
    nlayers = len(W)
    ntotal = 0
    for i in range(nlayers):
        ntotal += W[i].shape[0] * W[i].shape[1] + b[i].shape[0] * b[i].shape[1]
    v = np.zeros([ntotal, 1])
    vcount = 0
    for i in range(nlayers):
        vinc = W[i].shape[0] * W[i].shape[1]
        v[vcount:vcount + vinc] = np.reshape(W[i], (vinc, 1))
        vcount += vinc
    for i in range(nlayers):
        vinc = b[i].shape[0] * b[i].shape[1]
        v[vcount:vcount + vinc] = np.reshape(b[i], (vinc, 1))
        vcount += vinc

    return v


def vector_to_cell(v, W, b):
    """
    Reshapes an array from a 1-D vector to a dictionary-type structure
    """

    nlayers = len(W)
    vcount = 0

    for i in range(nlayers):
        vinc = W[i].shape[0] * W[i].shape[1]
        W[i] = np.reshape(v[vcount:vcount + vinc], (W[i].shape[0], W[i].shape[1]))
        vcount += vinc

    for i in range(nlayers):
        vinc = b[i].shape[0] * b[i].shape[1]
        b[i] = np.reshape(v[vcount:vcount + vinc], (b[i].shape[0], b[i].shape[1]))
        vcount += vinc

    return W, b
