import pandas as pd
import numpy as np
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.metrics import confusion_matrix

from seaborn import despine
import seaborn as sns
sns.set_style("ticks")
sns.set_context("talk")

from IPython.display import Image

import matplotlib.pyplot as plt
%matplotlib inline

# our activation function:

def sigmoid(x):
    """the sigmoid function
    """
    return 1.0/(1.0 + np.exp(-x)) 

# our loss: 

class cross_entropy_loss:
    def __init__(self):
        self.name = 'cross-entropy'
        
    def loss(self, y, y_pred, zerotol=1e-10):
        """the binary cross-entropy loss
        
        Args:
            y (array): labels for each instance (0 or 1)
            y_pred (array): predicted probabilty that
                each instance belongs to class 1
        """
        loss = -(y * np.log(y_pred + zerotol) + (1 - y) * np.log(1 - y_pred + zerotol))
        return loss
    
    def derivative_loss(self, y, y_pred):
        """the derivative of the cross-entropy 
        loss w.r.t. to sigmoid activation
        
        Args:
            y (array): labels for each instance (0 or 1)
            y_pred (array): predicted probabilty that
                each instance belongs to class 1
        """
        return y_pred - y

# the Perceptron:

class Perceptron:
    
    def __init__(self, n_in, activation=sigmoid, loss=cross_entropy_loss, b=None):
        """A simple Perceptron implementation.
        
        Args:
            n_in (int): number of input features for each instance
            activation (function): activation function of the Perceptron;
                takes individual values or array as input
            loss (class): loss function to use; 
                class is expected to include two functions:
                loss, derivative_loss (indicating the derivative
                of the loss w.r.t. to the output activation)
                see cross_entropy class above as an example
            b (float): bias term; if a value is specified, the
                bias term is fixed to this value. if not, 
                the bias will be estimated during training.
        """
        self.n_in = n_in
        self.w = np.random.uniform(-1,1,n_in)
        if b is None:
            self.b = np.random.uniform(-1,1,1)
            self.fit_b = True
        else:
            self.b = b
            self.fit_b = False
        self.activation = activation
        self.loss = loss().loss
        self.derivative_loss = loss().derivative_loss

    def predict(self, x):
        """Predict probability that each 
        instance of x (with smape n_instances x n_features)
        belongs to class 1
        
        Args:
            x (ndarray): input data (n_instances x n_features)
        """
        self.Z = np.dot(x, self.w) + self.b
        self.A = self.activation(self.Z)
        return self.A
    
    def gradient_descent_step(self, x, y, learning_rate):
        """A single gradient descent step.
        
        Args:
            x (ndarray): input data (n_instances x n_features)
            y (array): label of each instance (0 or 1)
            learning_rate (float): learning rate of gradient
                descent algorithm
        """
        # compute derivative of loss wrt Z
        dZ = self.derivative_loss(y, self.predict(x))
        dW = np.dot(dZ, x)
        # subtract average derivative from weights
        self.w -= learning_rate * 1.0/dW.shape[0] * dW
        if self.fit_b:
            self.b -= learning_rate * (1.0/x.shape[0] * np.sum(dZ))
            
    def train(self, x, y, batch_size=8, learning_rate=1, n_steps=100):
        """Iteratively train the Perceptron.
        At each iteration, the algorithm will draw 
        a random sample from x (with batch_size samples) 
        and use this sample to perform a gradient
        descent step.
        
        Args:
            x (ndarray): input data (n_instances x n_features)
            y (array): label of each instance (0 or 1)
            learning_rate (float): learning rate of gradient
                descent algorithm
            n_steps (int): number of iterations to perform
                during training
        """
        self.training_w = np.zeros((n_steps, self.n_in+1))
        self.training_loss = np.zeros(n_steps)
        for s in tqdm(range(n_steps)):
            # draw a random batch
            batch_idx = np.random.choice(x.shape[0], batch_size, replace=False)
            # compute and store mean loss
            self.training_loss[s] = np.mean(self.loss(y[batch_idx], self.predict(X[batch_idx])))
            # store current weights
            self.training_w[s,:self.n_in] = self.w
            self.training_w[s,-1] = self.b
            # perform gradient descent step
            self.gradient_descent_step(X[batch_idx], y[batch_idx], learning_rate)

# standardize predictions
p1_pred = (p1_pred - np.mean(p1_pred, axis=0)) / np.std(p1_pred, axis=0)
p2_pred = (p2_pred - np.mean(p2_pred, axis=0)) / np.std(p2_pred, axis=0)

y3 = np.logical_or(y1==1, y2==1).astype(np.int)

np.random.seed(1342)
# weights and bias of hidden layer:
W_hidden  = np.random.normal(size=(2,10))
b_hidden = np.random.normal(size=(1,10)) # one bias for each neuron

np.random.seed(473)
# weights and bias of output layer:
W_output  = np.random.normal(size=(10,1))
b_output = np.random.normal(size=(1,1)) # one bias for each neuron

# the sigmoid activation function

class sigmoid():
    """the sigmoid function
    """
    def __init__(self):
        self.name = 'sigmoid'
        
    def forward(self, x):
        """forward pass with the sigmoid function
        """
        return 1./(1.+np.exp(-x))
    
    def derivative(self, a):
        """the derviative phi'(a) of the sigmoid function
        """
        return a*(1-a)

# the neural network

class NeuralNetwork:
    """A simple Neural Network implementation.
    
    Args: 
        n_in (int): number of input neurons
        n_out (int): number of output neurons
        n_hidden (array of ints): number of neurons for each layer:
            an array of [10, 10] would create two hidden 
            layers with 10 neurons each
        activations (array of classes): activation for each hidden layer
            as well as the output layer; This should contain one activation
            function for each of the hidden layers (as specified by 
            n_hidden) as well as one for the output layer (this should be 
            the last entry of the array). If None, all activations are
            set to the sigmoid function.
            Importantly, each activation needs to be a
            class with two functions for the forward pass and
            the derivative w.r.t. its activations 
            (see the sigmoid implementation above)
        seed (int): random seed to allow for reproducibiliy
            
    """
            
    def __init__(self, n_in, n_out, n_hidden=[10], activations=None, seed=123):
        
        # init
        self.n_in = n_in
        self.n_out = n_out
        self.n_hidden = np.append(np.array(n_hidden), n_out).astype(np.int) # add output layer
        self.n_layers = self.n_hidden.size
        self.training_loss = [] # we will use this to collect the loss values during gradient descent
        self.seed = seed
        
        # initialize weights & biases
        self.parms = {} # a dictionary containing all network parameters    
        np.random.seed(self.seed)
        n_in = int(self.n_in)
        for layer, n_out in enumerate(self.n_hidden):
            self.parms["W{}".format(layer)] = self._initialize_weight([n_in, n_out])
            self.parms["b{}".format(layer)] = self._initialize_weight([1, n_out])
            # also create empty storage for neuron activations
            self.parms["A{}".format(layer)] = np.zeros((n_in, n_out)) 
            n_in = int(n_out)
        
        # set layer activations
        if activations is None:
            # if no activations are given, use sigmoid
            self.activations = [sigmoid()] * self.n_layers
        else:
            if len(activations) != self.n_layers:
                raise ValueError('/!\ Number of activations does not match number of layers.')
            self.activations = [a() for a in activations]
             
                
    def _initialize_weight(self, size, bound=1):
        """Draw weights randomly from 
        a uniform distribution between -1 and 1
        """
        return np.random.uniform(-1, 1, size=size)
    
    
    def forward(self, X):
        """Forward pass of the neural network,
        given a dataset X
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
                
        Returns:
            predicted probability (n_samples x n_output_neurons)
        """
        self.parms['X'] = X
        for layer in range(self.n_layers):
            Z = X.dot(self.parms['W{}'.format(layer)]) + self.parms['b{}'.format(layer)]
            A = self.activations[layer].forward(Z)
            X = A
            self.parms["A{}".format(layer)] = A
        self.y_pred = A
        return A
        
        
    def predict(self, X):
        """Predict classes for each instance
         of an input dataset X
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
                
        Returns:
            predicted class for each instance
        """
        y_pred = self.forward(X)
        if self.n_out == 1:
            return np.array(y_pred > 0.5).astype(np.int).ravel()
        else:
            return np.argmax(y_pred, axis=1).ravel()
    
    
    def loss(self, X, y, zerotol=1e-10):
        """Compute mean cross entropy loss,
        over the samples of an input dataset (X) and a set 
        of corresponding class labels (y).
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
            y (ndarray): class label for each instance. 
                Importantly, if n_out > 1, y needs to be 
                one-hot encoded: such that it is of
                shape (n_samples x n_out)
        
        Returns:
            Average cross entropy loss"""
        if self.activations[-1].name not in ['sigmoid', 'softmax']:
            raise ValueError('loss function only valid for sigmoid or softmax output activations.')
        y_pred = self.forward(X)
        # binary cross-entropy loss
        if self.n_out == 1:
            loss = -y * np.log(y_pred + zerotol) 
            loss -= (1 - y) * np.log(1 - y_pred + zerotol)
        # multi-class extension
        else:
            loss = -np.sum(y * np.log(y_pred), axis=1)
        return 1.0/X.shape[0] * np.sum(loss, axis=0)
    
    
    def derivative_loss(self, x, y):
        """Derivative of cross entropy loss"""
        if self.activations[-1].name not in ['sigmoid', 'softmax']:
            raise ValueError('derivative of loss only valid for sigmoid / softmax output activations.')
        y_pred = self.forward(x)
        return y_pred - y
        
        
    def backward(self, X, y, lr=0.1):
        """Compute gradient descent backward pass,
        updating all network weights, given an
        input dataset (X) and a set of corresponding 
        class labels (y).
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
            y (ndarray): class label for each instance. 
                Importantly, if n_out > 1, y needs to be 
                one-hot encoded: such that it is of
                shape (n_samples x n_out)
            lr (float): learning rate for weight update
        """
        y_pred = self.forward(X)
        self.training_loss.append(self.loss(X, y))
        dZ = self.derivative_loss(X, y)
        for layer in range(1, self.n_layers)[::-1]: # iterate hidden layers backward
            # dL/dW and dL/db are computed as averages over the batch
            # dL/dW is computed by multiplying dZ with the input to the layer 
            # (here, A from the next lower layer)
            dW = 1.0/dZ.shape[0] * self.parms["A{}".format(layer-1)].T.dot(dZ)
            db = 1.0/dZ.shape[0] * np.sum(dZ, axis=0, keepdims=True)
            self.parms["W{}".format(layer)] -= self.lr * dW
            self.parms["b{}".format(layer)] -= self.lr * db
            # compute dZ for next lower layer
            dA = dZ.dot(self.parms["W{}".format(layer)].T)
            dZ = dA * self.activations[layer-1].derivative(self.parms["A{}".format(layer-1)])
        # input layer
        dW = 1.0/dZ.shape[0] * X.T.dot(dZ)
        db = 1.0/dZ.shape[0] * np.sum(dZ, axis=0, keepdims=True)
        self.parms["W0"] -= self.lr * dW
        self.parms["b0"] -= self.lr * db
          
            
    def fit(self, X, y, lr=0.1, n_steps=10000, batch_size=32, verbose=False):
        """Fit the neural network to an input dataset.
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
            y (ndarray): class label for each instance. 
                Importantly, if n_out > 1, y needs to be 
                one-hot encoded: such that it is of
                shape (n_samples x n_out)
            lr (float): learning rate for weight update
            n_steps (int): number of gradient descent steps
                to perform
            batch_size (int): number of random samples 
                drawn from the dataset at each gradient descent iteration
            verbose (bool): whether or not to give a printed update
                about the state of the training"""
        # make sure y is right shape
        if self.n_out == 1 and np.ndim(y) < 2:
            y = y.reshape(-1,1)  
        assert np.ndim(y) == 2, '/!\ y needs to be 2 dimenional'
        # make sure data and model match
        if y.shape[1]>1:
            if np.unique(y.argmax(axis=1)).size != self.n_out:
                raise ValueError('/!\ n_out does not match number of labels in y (dim-1)')
        if verbose:
            print('Beginning training for {} batches ({} samples/batch):'.format(n_steps, batch_size))
        self.lr = lr
        iterator = range(n_steps)
        for step in (tqdm(iterator) if verbose else iterator):
            batch_i = np.random.choice(X.shape[0], batch_size, replace=False)
            self.forward(X[batch_i])
            self.backward(X[batch_i], y[batch_i], lr=self.lr)
         
        
    def plot_training_stats(self, X, y, train_idx, test_idx, target_names=None):
        """Plot the training loss as well as a confusion matrix
        for the training and test datasets.
        
        Args: 
            X (ndarray): input data with one instance 
                per row and one column for each feature
            y (array): class label for each instance. 
            train_idx (array): index indicating the
                position of the training instances in X and y
            test_idx (array): index indicating the
                position of the test instances in X and y
            target_names (array of strings): names of 
                target classes in y
                
        Returns:
            Matrplotlib figure and axes
        """
        fig, axs = plt.subplots(1,3,figsize=(20,6))
        # plot training loss
        axs[0].set_title('Training loss')
        axs[0].set_ylabel('Loss')
        axs[0].set_xlabel('Training step')
        axs[0].plot(self.training_loss, color='k')
        despine(ax=axs[0])
        # plot confusion matrix for training and test data
        for i, (label, idx) in enumerate(zip(['Training', 'Test'], [train_idx, test_idx])):
            y_pred = self.predict(X[idx])
            acc = np.mean(y_pred == y[idx])
            axs[1+i].set_title('{} data\nMean Acc.: {}%'.format(label, np.round(acc*100, 2)))
            conf_mat = confusion_matrix(y[idx], y_pred, normalize='true')
            sns.heatmap(conf_mat, annot=True, ax=axs[1+i], vmin=0, vmax=1,
                        annot_kws={'fontsize': 100./self.n_out}) # this might be too large for few classes
            if target_names is not None:
                axs[1+i].set_xticklabels(target_names)
                axs[1+i].set_yticklabels(target_names)
        fig.tight_layout()
        return fig, axs

# three common activation functions:

class relu():
    def __init__(self):
        self.name = 'relu'
        
    def forward(self, x):
        return np.maximum(x, 0)
    
    def derivative(self, a):
        d = np.zeros_like(a)
        d[a<0] = 0
        d[a>0] = 1
        return d
    
    
class tanh():
    def __init__(self):
        self.name = 'tanh'
        
    def forward(self, x):
        return (np.exp(x)-np.exp(-x))/(np.exp(x)+np.exp(-x))
    
    def derivative(self, a):
        return 1-a**2
    
    
class softmax():
    def __init__(self):
        self.name = 'softmax'
        
    def forward(self, x):
        e_x = np.exp(x)
        return e_x / e_x.sum(axis=1, keepdims=True)
    
    def derivative(self, a):
        SM = a.reshape((-1,1))
        return (np.diagflat(a) - np.dot(SM, SM.T)).sum(axis=1)

from sklearn.datasets import load_digits

data = load_digits()

# create X & y data
X = data['data']
y = data['target']

# create a one-hot version of y
y_onehot = np.zeros((y.size, np.unique(y).size))
y_onehot[np.arange(y.size), y] = 1
assert np.all(np.argmax(y_onehot, axis=1) == y), '/!\ Error in one-hot encoding'

np.random.seed(1827)
# train / test split
train_idx = np.random.choice(X.shape[0], int(0.7 * X.shape[0]), replace=False)
test_idx = np.array([i for i in np.arange(X.shape[0]) if i not in train_idx])
assert np.all([i not in train_idx for i in test_idx]), '/!\ Overlapping training and test datasets.'