from __future__ import print_function

import numpy as np
import matplotlib.pyplot as plt


class TwoLayerNet(object):
    """
    A two-layer fully-connected neural network. The net has an input dimension
    of N, a hidden layer dimension of H, and performs classification over C 
    classes.
    We train the network with a softmax loss function and L2 regularization on
    the weight matrices. The network uses a ReLU nonlinearity after the first
    fully connected layer.
    In other words, the network has the following architecture:
    input - fully connected layer - ReLU - fully connected layer - softmax
    The outputs of the second fully-connected layer are the scores for each
    class.
    """

    def __init__(self, input_size, hidden_size, output_size, std=1e-4):
        """
        Initialize the model. Weights are initialized to small random values
        and biases are initialized to zero. Weights and biases are stored in
        the variable self.params, which is a dictionary with the following keys
        W1: First layer weights; has shape (D, H)
        b1: First layer biases; has shape (H,)
        W2: Second layer weights; has shape (H, C)
        b2: Second layer biases; has shape (C,)
        Inputs:
        - input_size: The dimension D of the input data.
        - hidden_size: The number of neurons H in the hidden layer.
        - output_size: The number of classes C.
        """
        self.params = {}
        self.params['W1'] = std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

    def loss(self, X, y=None, reg=0.0):
        """
        Compute the loss and gradients for a two layer fully connected neural
        network.
        Inputs:
        - X: Input data of shape (N, D). Each X[i] is a training sample.
        - y: Vector of training labels. y[i] is the label for X[i], and each
          y[i] is an integer in the range 0 <= y[i] < C. This parameter is
          optional; if it is not passed then we only return scores, and if it
          is passed then we instead return the loss and gradients.
        - reg: Regularization strength.
        Returns:
        If y is None, return a matrix scores of shape (N, C) where scores[i, c]
        is the score for class c on input X[i].
        If y is not None, instead return a tuple of:
        - loss: Loss (data loss and regularization loss) for this batch of
          training samples.
        - grads: Dictionary mapping parameter names to gradients of those
          parameters with respect to the loss function; has the same keys as
          self.params.
        """
        # Unpack variables from the params dictionary
        W1, b1 = self.params['W1'], self.params['b1']
        W2, b2 = self.params['W2'], self.params['b2']
        N, D = X.shape

        # Compute the forward pass
        scores = None
        #######################################################################
        # TODO: Perform the forward pass, computing the class scores for the  #
        # input. Store the result in the scores variable, which should be an  #
        # array of shape (N, C).                                              #
        #######################################################################

        scores1 = X.dot(W1) + b1  # FC1
        X2 = np.maximum(0, scores1)  # ReLU FC1

        scores = X2.dot(W2) + b2  # FC2

        #######################################################################
        #                              END OF YOUR CODE                       #
        #######################################################################

        # If the targets are not given then jump out, we're done
        if y is None:
            return scores

        scores -= np.max(scores)  # Fix Number instability
        scores_exp = np.exp(scores)
        probs = scores_exp / np.sum(scores_exp, axis=1, keepdims=True)

        # Compute the loss
        loss = None
        #######################################################################
        # TODO: Finish the forward pass, and compute the loss. This should    #
        # include both the data loss and L2 regularization for W1 and W2.     #
        # Store the result in the variable loss, which should be a scalar. Use#
        # the Softmax classifier loss.                                        #
        #######################################################################
        
        correct_probs = -np.log(probs[np.arange(N), y])

        # L_i = -log(e^correct_score/sum(e^scores))) = -log(correct_probs)
        loss = np.sum(correct_probs)
        loss /= N
        # L2 regularization WRT W1 and W2
        loss += reg * (np.sum(W1 * W1) + np.sum(W2 * W2))

        #######################################################################
        #                              END OF YOUR CODE                       #
        #######################################################################

        # Backward pass: compute gradients
        grads = {}
        #############################################################################
        # TODO: Compute the backward pass, computing the derivatives of the weights #
        # and biases. Store the results in the grads dictionary. For example,       #
        # grads['W1'] should store the gradient on W1, and be a matrix of same size #
        #############################################################################

        # gradient of loss_i WRT scores_k
        # dL_i/ds_k = probs_k-1(y_i == k)
        # this means the gradient is the score for "other" classes and score-1
        # for the target class
        d_scores = probs.copy()
        d_scores[np.arange(N), y] -= 1
        d_scores /= N

        # W2 were multiplied with X2, by chain rule and multiplication
        # derivative, WRT W2 we need to multiply downstream derivative by X2
        d_W2 = X2.T.dot(d_scores)
        # b2 was added, so it's d is 1 but we must multiply it with chain rule
        # (downstream), in this case d_scores
        d_b2 = np.sum(d_scores, axis=0)

        # W1 is upstream of X2, so we continue this way
        d_X2 = d_scores.dot(W2.T)
        # ReLU derivative is 1 for > 0, else 0
        d_scores1 = d_X2 * (scores1 > 0)
        d_W1 = X.T.dot(d_scores1)

        # b1 gradient
        d_b1 = d_scores1.sum(axis=0)

        # regularization gradient (reg*W2^2)
        d_W2 += reg * 2 * W2
        d_W1 += reg * 2 * W1

        grads['W1'] = d_W1
        grads['b1'] = d_b1
        grads['W2'] = d_W2
        grads['b2'] = d_b2
        #######################################################################
        #                              END OF YOUR CODE                       #
        #######################################################################

        return loss, grads

    def train(self, X, y, X_val, y_val,
              learning_rate=1e-3, learning_rate_decay=0.95,
              reg=5e-6, num_iters=100,
              batch_size=200, verbose=False):
        """
        Train this neural network using stochastic gradient descent.
        Inputs:
        - X: A numpy array of shape (N, D) giving training data.
        - y: A numpy array f shape (N,) giving training labels; y[i] = c means
          that X[i] has label c, where 0 <= c < C.
        - X_val: A numpy array of shape (N_val, D) giving validation data.
        - y_val: A numpy array of shape (N_val,) giving validation labels.
        - learning_rate: Scalar giving learning rate for optimization.
        - learning_rate_decay: Scalar giving factor used to decay the learning
          rate after each epoch.
        - reg: Scalar giving regularization strength.
        - num_iters: Number of steps to take when optimizing.
        - batch_size: Number of training examples to use per step.
        - verbose: boolean; if true print progress during optimization.
        """
        num_train = X.shape[0]
        iterations_per_epoch = max(num_train / batch_size, 1)

        # Use SGD to optimize the parameters in self.model
        loss_history = []
        train_acc_history = []
        val_acc_history = []

        for it in range(num_iters):
            X_batch = None
            y_batch = None

            ###################################################################
            # TODO: Create a random minibatch of training data and labels,    #
            # storing  them in X_batch and y_batch respectively.              #
            ###################################################################

            # random indexes to sample training data/labels
            sample_idx = np.random.choice(num_train, batch_size, replace=True)
            X_batch = X[sample_idx]
            y_batch = y[sample_idx]

            ###################################################################
            #                             END OF YOUR CODE                    #
            ###################################################################

            # Compute loss and gradients using the current minibatch
            loss, grads = self.loss(X_batch, y=y_batch, reg=reg)
            loss_history.append(loss)

            ###################################################################
            # TODO: Use the gradients in the grads dictionary to update the   #
            # parameters of the network (stored in the dictionary self.params)#
            # using stochastic gradient descent. You'll need to use the       #
            # gradients stored in the grads dictionary defined above.         #
            ###################################################################

            # For each weight in network parameters, update it with the
            # corresponding calculated gradient
            for key in self.params:
                self.params[key] -= learning_rate * grads[key]

            ###################################################################
            #                             END OF YOUR CODE                    #
            ###################################################################

            if verbose and it % 100 == 0:
                print('iteration %d / %d: loss %f' % (it, num_iters, loss))

            # Every epoch, check train and val accuracy and decay learning rate
            if it % iterations_per_epoch == 0:
                # Check accuracy
                train_acc = (self.predict(X_batch) == y_batch).mean()
                val_acc = (self.predict(X_val) == y_val).mean()
                train_acc_history.append(train_acc)
                val_acc_history.append(val_acc)

                # Decay learning rate
                learning_rate *= learning_rate_decay

        return {
            'loss_history': loss_history,
            'train_acc_history': train_acc_history,
            'val_acc_history': val_acc_history,
        }

    def predict(self, X):
        """
        Use the trained weights of this two-layer network to predict labels for
        data points. For each data point we predict scores for each of the C
        classes, and assign each data point to the class with the highest score
        Inputs:
        - X: A numpy array of shape (N, D) giving N D-dimensional data points
          to classify.
        Returns:
        - y_pred: A numpy array of shape (N,) giving predicted labels for each
          of the elements of X. For all i, y_pred[i] = c means that X[i] is
          predicted to have class c, where 0 <= c < C.
        """
        y_pred = None

        #######################################################################
        # TODO: Implement this function; it should be VERY simple!            #
        #######################################################################
        y_pred = np.argmax(self.loss(X), axis=1)
        #######################################################################
        #                              END OF YOUR CODE                       #
        #######################################################################

        return y_pred
[((37, 28, 37, 49), 'numpy.zeros', 'np.zeros', ({(37, 37, 37, 48): 'hidden_size'}, {}), '(hidden_size)', True, 'import numpy as np\n'), ((39, 28, 39, 49), 'numpy.zeros', 'np.zeros', ({(39, 37, 39, 48): 'output_size'}, {}), '(output_size)', True, 'import numpy as np\n'), ((76, 13, 76, 35), 'numpy.maximum', 'np.maximum', ({(76, 24, 76, 25): '0', (76, 27, 76, 34): 'scores1'}, {}), '(0, scores1)', True, 'import numpy as np\n'), ((88, 18, 88, 32), 'numpy.max', 'np.max', ({(88, 25, 88, 31): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((89, 21, 89, 35), 'numpy.exp', 'np.exp', ({(89, 28, 89, 34): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((104, 15, 104, 36), 'numpy.sum', 'np.sum', ({(104, 22, 104, 35): 'correct_probs'}, {}), '(correct_probs)', True, 'import numpy as np\n'), ((134, 15, 134, 39), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((36, 34, 36, 74), 'numpy.random.randn', 'np.random.randn', ({(36, 50, 36, 60): 'input_size', (36, 62, 36, 73): 'hidden_size'}, {}), '(input_size, hidden_size)', True, 'import numpy as np\n'), ((38, 34, 38, 75), 'numpy.random.randn', 'np.random.randn', ({(38, 50, 38, 61): 'hidden_size', (38, 63, 38, 74): 'output_size'}, {}), '(hidden_size, output_size)', True, 'import numpy as np\n'), ((90, 29, 90, 70), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((197, 25, 197, 78), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((107, 23, 107, 38), 'numpy.sum', 'np.sum', ({(107, 30, 107, 37): '(W1 * W1)'}, {}), '(W1 * W1)', True, 'import numpy as np\n'), ((107, 41, 107, 56), 'numpy.sum', 'np.sum', ({(107, 48, 107, 55): '(W2 * W2)'}, {}), '(W2 * W2)', True, 'import numpy as np\n'), ((126, 17, 126, 29), 'numpy.arange', 'np.arange', ({(126, 27, 126, 28): 'N'}, {}), '(N)', True, 'import numpy as np\n'), ((101, 38, 101, 50), 'numpy.arange', 'np.arange', ({(101, 48, 101, 49): 'N'}, {}), '(N)', True, 'import numpy as np\n')]