# GRADED CODE: multi-class classification
### START CODE HERE ###
def random_mini_batches(X, Y, mini_batch_size = 64):
    """
    Creates a list of random minibatches from (X, Y)

    Arguments:
    X -- input data, of shape (n, f^{0})
    Y -- true "label" vector, of shape (n, C)
    mini_batch_size -- size of the mini-batches, integer

    Returns:
    mini_batches -- list of synchronous (mini_batch_X, mini_batch_Y)
    """

    m = X.shape[0]                  # number of training examples
    mini_batches = []

    # Step 1: Shuffle (X, Y)
    permutation = list(np.random.permutation(m))
    shuffled_X = X[permutation, :]
    shuffled_Y = Y[permutation, :]

    # Step 2 - Partition (shuffled_X, shuffled_Y).
    # Cases with a complete mini batch size only i.e each of 64 examples.
    num_complete_minibatches = math.floor(m / mini_batch_size) # number of mini batches of size mini_batch_size in your partitionning
    for k in range(0, num_complete_minibatches):
        # (approx. 2 lines)
        mini_batch_X = shuffled_X[k * mini_batch_size : (k + 1) * mini_batch_size, :]
        mini_batch_Y = shuffled_Y[k * mini_batch_size : (k + 1) * mini_batch_size, :]
        mini_batch = (mini_batch_X, mini_batch_Y)
        mini_batches.append(mini_batch)

    # For handling the end case (last mini-batch < mini_batch_size i.e less than 64)
    if m % mini_batch_size != 0:
        #(approx. 2 lines)
        mini_batch_X = shuffled_X[num_complete_minibatches * mini_batch_size : , :]
        mini_batch_Y = shuffled_Y[num_complete_minibatches * mini_batch_size : , :]
        mini_batch = (mini_batch_X, mini_batch_Y)
        mini_batches.append(mini_batch)

    return mini_batches

classes = 10
layers_dims = [x_train.shape[1], 128, 64, classes]
activation_fn = ["relu", "relu", "softmax"]
learning_rate = 0.01
num_iterations = 100
batch_size = 64
losses = []                         # keep track of loss
print_loss = True
print_freq = 10
loss_function = 'cross_entropy'
gamma = None
alpha = None
model = Model(layers_dims, activation_fn, loss_function, alpha=alpha, gamma=gamma)

# Loop (gradient descent)
for i in range(0, num_iterations):
    mini_batches = random_mini_batches(x_train, y_train, batch_size)
    loss = 0
    for batch in mini_batches:
        x_batch, y_batch = batch

        # forward
        #x_batch_flattened = x_batch.reshape(x_batch.shape[0], -1)
        AL = model.forward(x_batch)

        # compute loss
        if loss_function == 'cross_entropy':
            loss += compute_CCE_loss(AL, y_batch)
        elif loss_function == 'focal_loss':
            loss += compute_focal_loss(AL, y_batch, gamma, alpha)

        # backward
        dA_prev = model.backward(AL, y_batch)
        # update
        model.update(learning_rate)

    loss /= len(mini_batches)
    losses.append(loss)
    if print_loss and i % print_freq == 0:
        print ("Loss after iteration %i: %f" %(i, loss))


# plot the loss
plt.figure(figsize=(4, 2))
plt.plot(np.squeeze(losses))
plt.ylabel('loss')
plt.xlabel('iterations (per hundreds)')
plt.title("Learning rate =" + str(learning_rate))
plt.show()
### END CODE HERE ###