import numpy as np
from sigmoid import *

def predict(X, Theta, num_layers=2):
    """
    Predict the label of an input given a tarined neural network

    Parameters:
    -----------
    X : 2D-array (m_examples, n_features)
        input for the hypothesis to predict.

    Theta :
        Array of weights between layers
    """
    m_examples, n_features = X.shape

    # Input Layer
    a1 = np.r_[np.ones((1, m_examples)), X.T]

    # Hidden Layer
    a2 = sigmoid(Theta[0].dot(a1))
    a2 = np.r_[np.ones((1,m_examples)), a2]

    # Output Layer
    a3 = sigmoid(Theta[1].dot(a2))

    # Return predictions
    return np.argmax(a3, axis=0)