## ~~~~~~~~~~~~~~~~~~
#  Deep Willy Network
## ~~~~~~~~~~~~~~~~~~

import numpy as np, json, sys, os
sys.path.append(os.path.dirname(__file__))

from willies import *


class DeepWilly(object):
    
    _willy_classes = {'connected': ConnectedWilly,
                      'dropout': DropoutWilly,
                      'convolutional': ConvolutionalWilly,
                      'pooling': PoolingWilly,
                      'stacking': StackingWilly}
    
    def __init__(self, cost_func = 'cross entropy'):
        
        self.willies = []
        self.n_willies = 0
        
        self.cost_func = cost_func
        if cost_func == 'cross entropy': 
            self.cost  = lambda a, y: -np.mean(y * np.log(a + (a == 0)) + (1-y) * np.log(1-a + (a == 1)))
            self.dcost = lambda a, y: (1/a.shape[0]) * (a - y) / (a * (1 - a) + 1e-99)
        
        elif cost_func == 'quadratic':
            self.cost  = lambda a, y: 0.5 * np.mean((a - y)**2)
            self.dcost = lambda a, y: (1/a.shape[0]) * (a - y)
        
        else:
            raise Exception('Unsupported cost function: ' + cost_func)
    
    def add(self, willy):
        ''' Add a willy to the network and set it up based on the previous willy's output (if not the first one). '''
        
        if len(self.willies) == 0:
            if not willy.is_set_up:
                raise Exception('Input shape or number must be provided to first Willy.')
                
        else:
            willy.set_up(self.willies[-1].out_shape)
            
        self.willies.append(willy)
        self.n_willies += 1
        
    def forward_prop(self, X):
        ''' Forward propagates X through the willies. '''
        
        for willy in self.willies:
            X = willy.forward_prop(X)
        
        return X
    
    def backward_prop(self, X, y,
                      learn_rate, mom_rate, reg_rate):
        ''' Backward propagates errors while simultaneously updating weights in each layer. '''
        
        # Output layer dC/da
        batch_size = X.shape[0]
        dA = self.dcost(self.willies[-1].A, y)
        
        # Backpropagate, updating weights
        XA = [X] + [willy.A for willy in self.willies[:-1]]
        XA = XA[::-1]
        
        for w, willy in enumerate(reversed(self.willies)):
            willy.update_weights(dA, XA[w],
                                 learn_rate, mom_rate, reg_rate)
            dA = willy.backward_prop()
    
    @staticmethod
    def get_batches(X, y, batch_size):
        ''' Shuffles training data and gets random batch of desired size. '''
        
        if batch_size == -1 or batch_size >= X.shape[0]:
            return [X], [y]
        
        # Shuffle data
        shuffled_indices = np.random.permutation(len(X))
        shuffled_X = X[shuffled_indices]
        shuffled_y = y[shuffled_indices]

        # Get batch of desired size
        X_batches = []
        y_batches = []
        for i in range(X.shape[0]//batch_size):
            X_batches.append(shuffled_X[int(batch_size*i):int(batch_size*(i+1))])
            y_batches.append(shuffled_y[int(batch_size*i):int(batch_size*(i+1))])
        
        return X_batches, y_batches
    
    def train(self, X, y, 
              num_iterations, batch_size, 
              learn_rate, reg_rate, mom_rate = 0,
              verbose = False):
        ''' Builds network, trains using given data and training parameters. '''
        
        # Change dtypes
        X = X.astype(np.float32)
        y = y.astype(np.float32)
        
        # Initialise momenta to 0
        for willy in self.willies:
            willy.reset_momenta()
        
        # Train network
        for iteration in range(num_iterations):
            
            # Get batches:
            X_batches, y_batches = self.get_batches(X, y, batch_size)
            
            for batchX, batchy in zip(X_batches, y_batches):
                
                # Forward propagate
                self.forward_prop(batchX)

                # Backward propagate & update weights
                self.backward_prop(batchX, batchy,
                                   learn_rate, mom_rate, reg_rate)

            # Print progress
            if verbose:
                if iteration % verbose == 0:
                    print("Training cost on last batch: ", self.cost(self.willies[-1].A, batchy))
    
    def predict(self, X, pred_type = 'as is'):
        
        self.yhat = self.forward_prop(X)
        
        if pred_type == 'binary':
            self.yhat = 1 * (self.yhat > 0.5)
            
        elif pred_type == 'argmax':
            self.yhat = np.argmax(self.yhat, axis = 1).reshape(-1, 1)
        
        else:
            assert pred_type == 'as is', \
                "Provided argument pred_type (" + pred_type + ") not supported."
        
        return self.yhat
    
    def accuracy(self, X, y, pred_type = 'as is'):
        ''' Gets accuracy of predictions. '''

        return np.mean(self.predict(X, pred_type) == y)
    
    def save(self, filename):
        ''' Saves deep willy to file \filename. '''
        
        willy_data = []
        for willy in self.willies:
            willy_data.append(willy.save())
        
        data = {'cost': self.cost_func,
                'willies': willy_data}
        
        file = open(filename, "w")
        json.dump(data, file)
        file.close()
    
    @classmethod
    def load(cls, filename):
        ''' Loads deep willy from file \filename. '''
        
        file = open(filename, "r")
        data = json.load(file)
        file.close()
        
        deep_willy = cls(data['cost_func'])
        for willy in data['willies']:
            willy_class = DeepWilly._willy_classes[willy['willy']]
            deep_willy.add(willy_class.load(willy))
            
        return deep_willy
    
    def copy(self):
        ''' Replicates this deep willy using its attributes. '''
        
        copy = DeepWilly(cost_func = self.cost_func)
        
        for willy in self.willies:
            copy.willies.append(willy.copy())
            copy.n_willies += 1
        
        return copy

