# GRADED CODE: binary classification
### START CODE HERE ###

loss_function = 'cross_entropy'
if loss_function == 'cross_entropy':
    layers_dims = [x_train.shape[1], 64, 1]
    activation_fn = ['relu', 'sigmoid']
    gamma = None        # you can leave this as it is
    alpha = None        # you can leave this as it is
    y_train_processed = y_train.reshape(-1, 1)
    y_val_processed = y_val.reshape(-1, 1)
    assert y_train_processed.shape[-1] == 1, "see the 'Note' in the Basic implementation section"
    assert y_val_processed.shape[-1] == 1, "see the 'Note' in the Basic implementation section"
elif loss_function == 'focal_loss':
    layers_dims = [x_train.shape[1], 64, 2]
    activation_fn = ['relu', 'softmax']
    gamma = 2.0
    alpha = None
    y_train_processed = np.hstack([1 - y_train, y_train])
    y_val_processed = np.hstack([1 - y_val, y_val])
    assert y_train_processed.shape[-1] == 2, "see the 'Note' in the Basic implementation section"
    assert y_val_processed.shape[-1] == 2, "see the 'Note' in the Basic implementation section"

learning_rate = 0.01
num_iterations = 10000
print_loss = True
print_freq = 1000
classes = 2
losses = []                         # keep track of loss
model = Model(layers_dims, activation_fn, loss_function, alpha, gamma)

# Loop (batch gradient descent)
for i in range(0, num_iterations):
    # forward
    AL = model.forward(x_train)

    # compute loss
    if loss_function == 'cross_entropy':
        loss = compute_BCE_loss(AL, y_train_processed)
    elif loss_function == 'focal_loss':
        loss =  compute_focal_loss(AL, y_train_processed, alpha, gamma)

    # backward
    dA_prev = model.backward(AL, y_train_processed)

    # update
    model.update(learning_rate)


    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=(6, 3))
plt.plot(np.squeeze(losses))
plt.ylabel('loss')
plt.xlabel(f'iterations (per {print_freq})')
plt.title("Learning rate =" + str(learning_rate))
plt.show()
### END CODE HERE ###