import numpy as np
import scipy.special as sp

class DeepGate:
    
    # Parameterize the network.
    def __init__(self, input_nodes, hidden_nodes, output_nodes, learning_rate):

        # Set the number of input nodes, hidden nodes per layer, output nodes,
        # and the learning rate.
        self.input_nodes = input_nodes
        self.hidden_nodes = hidden_nodes
        self.output_nodes = output_nodes
        self.learning_rate = learning_rate
        
        # Compute the depth and concatenate the input layer, hidden layer(s),
        # and the output layer.
        if type(self.hidden_nodes) == list:
            self.layers = [self.input_nodes] + self.hidden_nodes + \
                          [self.output_nodes]
        else:
            self.layers = [self.input_nodes] + [self.hidden_nodes] + \
                          [self.output_nodes]
        self.depth = len(self.layers)        
        
        # Generate the weight matrices.
        self.weights = []
        for layer in range(1, self.depth):
            mean = 0.0
            variance = np.power(self.layers[layer], -0.5)
            size = [self.layers[layer], self.layers[layer-1]]
            self.weights += [np.random.normal(mean, variance, size)]

        # Set the activation function(s) to use.
        self.activationFunction = lambda x: sp.expit(x)
        
    # Train the network.
    def train(self, input_layer, target):
        
        # Convert the input and target into 2-D arrays.
        input_layer = np.array(input_layer, ndmin=2).T
        target = np.array(target, ndmin=2).T
        
        # Forward pass the inputs at the input layer through the network.
        outputs = [input_layer]
        for layer in range(self.depth-1):
            temp = np.dot(self.weights[layer], outputs[layer])
            outputs += [self.activationFunction(temp)]
        output_layer = outputs[-1]
        
        # Compute the errors at the output layer.
        error = []
        error += [target - output_layer]
        
        # Backpropagate the errors at the output layer through the network.
        for layer in range(self.depth-2, 0, -1):
            error += [np.dot(self.weights[layer].T, 
                                error[self.depth-2-layer])]
        error = error[::-1]
        
        # Update the weights.  
        for layer in range(self.depth-2, -1, -1):
            temp = np.dot((error[layer] * \
                              outputs[layer+1] * \
                              (1-outputs[layer+1])), outputs[layer].T)
            self.weights[layer] += self.learning_rate * temp
    
    # Query the network.
    def query(self, input_layer):
        
        # Convert the input into a 2-D array.
        input_layer = np.array(input_layer, ndmin=2).T
        
        # Forward pass the inputs at the input layer through the network.
        outputs = [input_layer]
        for layer in range(self.depth-1):
            temp = np.dot(self.weights[layer], outputs[layer])
            outputs += [self.activationFunction(temp)]
        output_layer = outputs[-1]
        
        return output_layer
        