def small_test():
    num_inputs = 5
    input_size = 4
    hidden_size = 10
    ouptut_size = 3

    np.random.seed(0)
    model = TwoLayerNet(input_size=input_size, hidden_size=hidden_size, output_size=ouptut_size, reg=0.0, std=1e-1)

    np.random.seed(1)
    x = 10 * np.random.randn(num_inputs, input_size)
    y = np.array([0, 1, 2, 2, 1])

    # scores = model.compute_loss_and_gradients(x, None)
    return x, y, model

x, y, model = small_test()

scores = model.compute_loss_and_gradients(x, None)
loss, grads = model.compute_loss_and_gradients(x, y, reg=0.05)

correct_loss = 1.30378789133
correct_scores = np.asarray([
  [-0.81233741, -1.27654624, -0.70335995],
  [-0.17129677, -1.18803311, -0.47310444],
  [-0.51590475, -1.01354314, -0.8504215 ],
  [-0.15419291, -0.48629638, -0.52901952],
  [-0.00618733, -0.12435261, -0.15226949]])
# print("model.w1 =", model.params['w1'])
# print("model.w2 =", model.params['w2'])
print("correct_loss - loss =", np.sum(np.abs(correct_loss - loss)))
print("scores - correct_scores =", np.sum(np.abs(scores - correct_scores)))
print("xy =", np.sum(x), np.sum(y))
print("net.params =", {k.lower(): np.sum(v) for k, v in model.params.items()})
print("net.grads =", {k.lower(): np.sum(v) for k, v in grads.items()})
# print("params = ", {k: v.shape for k, v in params.items()})
# print("grads = ", {k: v.shape for k, v in grads.items()})

def small_train(x, y, batch_size=200):
    x_val = x
    y_val = y
    
    learning_rate=0.1
    learning_rate_decay=0.95
    num_iters=200
    reg=0.005
    num_train = x.shape[0]
    iterations_per_epoch = max(num_train / batch_size, 1)

    loss_history = []
    train_acc_history = []
    val_acc_history = []

    for it in range(num_iters):
        idx = np.random.choice(num_train, batch_size, replace=True)
        x_batch = x[idx]
        y_batch = y[idx]
        loss, grads = model.compute_loss_and_gradients(x_batch, y=y_batch)
        loss_history.append(loss)
        model.params['w2'] += - learning_rate * grads['w2']
        model.params['b2'] += - learning_rate * grads['b2']
        model.params['w1'] += - learning_rate * grads['w1']
        model.params['b1'] += - learning_rate * grads['b1']
        if it % 100 == 0:
            print('iteration %d / %d: loss %f' % (it, num_iters, loss))
        
        if it % iterations_per_epoch == 0:
            # Check accuracy
            train_acc = (model.predict(x_batch) == y_batch).mean()
            val_acc = (model.predict(x_val) == y_val).mean()
            train_acc_history.append(train_acc)
            val_acc_history.append(val_acc)
            learning_rate *= learning_rate_decay
    # finish the loop
    return {
        'loss_history': loss_history,
        'train_acc_history': train_acc_history,
        'val_acc_history': val_acc_history,
    }

stats = small_train(x, y, batch_size=200)

print('Final training loss: ', stats['loss_history'][-1])

# plot the loss history
plt.plot(stats['loss_history'])
plt.xlabel('iteration')
plt.ylabel('training loss')
plt.title('Training Loss history')
plt.show()