import numpy as np

from chapters.ex2.sigmoid import sigmoid
from sigmoidGradient import sigmoidGradient


def nnCostFunction(nn_params, input_layer_size, hidden_layer_size,
                   num_labels, X, y, Lambda):

    """ computes the cost and gradient of the neural network. The
        parameters for the neural network are "unrolled" into the vector
        nn_params and need to be converted back into the weight matrices.

        The returned parameter grad should be a "unrolled" vector of the
        partial derivatives of the neural network.
    """

    # Reshape nn_params back into the parameters Theta1 and Theta2, the weight matrices
    # for our 2 layer neural network
    # Obtain Theta1 and Theta2 back from nn_params
    Theta1 = np.reshape(nn_params[:hidden_layer_size * (input_layer_size + 1)],
                        (hidden_layer_size, input_layer_size + 1), order='F').copy()

    Theta2 = np.reshape(nn_params[hidden_layer_size * (input_layer_size + 1):],
                        (num_labels, (hidden_layer_size + 1)), order='F').copy()

    # Useful values
    m, n = X.shape
    num_labels, _ = Theta2.shape
    
    # Make computations for input layer
    a1 = np.column_stack((np.ones(m), X))
    
    # Make computations for hidden layer
    z2 = a1.dot(Theta1.T)
    a2 = sigmoid(z2)
    a2 = np.column_stack((np.ones(m), a2)) # add bias term
    
    # Make computations for output layer
    z3 = a2.dot(Theta2.T)
    h = sigmoid(z3)
    
    # Compute cost
    cost31 = np.log(h)
    cost30 = np.log(1 - h)
    regularization1 = Lambda/(2*m)*np.sum(np.array(Theta1)[:,1:]**2)
    regularization2 = Lambda/(2*m)*np.sum(np.array(Theta2)[:,1:]**2)
    regularization = regularization1 + regularization2
    
    Y = np.zeros((m, num_labels))
    for i in range(1, num_labels+1):
        Y[:,i-1] = np.array([1 if label == i else 0 for label in y])
    
    J = -1/m * np.sum( np.multiply(Y, cost31) + np.multiply(1.0-Y, cost30) ) + regularization
             
## =============================================================================
#    # Version with loops
#    D1 = np.zeros((input_layer_size+1, hidden_layer_size)).T
#    D2 = np.zeros((hidden_layer_size+1, num_labels)).T
#    
#    # Compute error
#    for t in range(0, m):
#        a1t = a1[t,:]
#        z2t = z2[t,:]
#        a2t = a2[t,:]
#        z3t = z3[t,:]
#        ht = h[t,:]
#        yt = Y[t,:]
#        d3t = ht - yt
#        
#        z2t = np.insert(z2t, 0, values=np.ones(1))
#        d2t = np.multiply(Theta2.T.dot(d3t[:,None]).T, sigmoidGradient(sigmoid(a2t)))
#        
#        D2 = D2 + np.outer(d3t.T,a2t)
#        D1 = D1 + np.dot(d2t[:,1:].T, a1t[:, None].T)
#        
#    D1 = D1 / m
#    D2 = D2 / m
    
# =============================================================================
    # Vectorized version
    
    d3 = h - Y 
    D2 = np.dot(d3.T,a2)
    
    d2 = np.multiply(d3.dot(Theta2[:,1:]), sigmoidGradient(z2))
    D1 = np.dot(d2.T,a1)
    
    D2 = 1/m*D2
    D1 = 1/m*D1
    
    #Regularization
    D1 = D1 + Lambda/m*np.concatenate((np.zeros((Theta1.shape[0],1)) ,Theta1[:,1:]), axis=1)
    D2 = D2 + Lambda/m*np.concatenate((np.zeros((Theta2.shape[0],1)) ,Theta2[:,1:]), axis=1)
    
    # ====================== YOUR CODE HERE ======================
    # Instructions: You should complete the code by working through the
    #               following parts.
    #
    # Part 1: Feedforward the neural network and return the cost in the
    #         variable J. After implementing Part 1, you can verify that your
    #         cost function computation is correct by verifying the cost
    #         computed in ex4.m
    #
    # Part 2: Implement the backpropagation algorithm to compute the gradients
    #         Theta1_grad and Theta2_grad. You should return the partial derivatives of
    #         the cost function with respect to Theta1 and Theta2 in Theta1_grad and
    #         Theta2_grad, respectively. After implementing Part 2, you can check
    #         that your implementation is correct by running checkNNGradients
    #
    #         Note: The vector y passed into the function is a vector of labels
    #               containing values from 1..K. You need to map this vector into a
    #               binary vector of 1's and 0's to be used with the neural network
    #               cost function.
    #
    #         Hint: We recommend implementing backpropagation using a for-loop
    #               over the training examples if you are implementing it for the
    #               first time.
    #
    # Part 3: Implement regularization with the cost function and gradients.
    #
    #         Hint: You can implement this around the code for
    #               backpropagation. That is, you can compute the gradients for
    #               the regularization separately and then add them to Theta1_grad
    #               and Theta2_grad from Part 2.
    #
    # =========================================================================

    # Unroll gradient
    grad = np.hstack((D1.T.ravel(), D2.T.ravel()))

    return J, grad
