def one_hot(y, num_classes):
    out = np.zeros((len(y), num_classes))
    out[np.arange(len(y)), y] = 1
    return out


if __name__ == "__main__":
    X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
    y = np.array([0, 1, 1, 0])  # XOR
    y_onehot = one_hot(y, 2)

    net = TwoLayerNet(d_in=2, d_hidden=8, d_out=2, seed=0)
    for epoch in range(3000):
        net.forward(X)
        net.backward(X, y_onehot, lr=0.5)
        if epoch % 500 == 0:
            print(f"epoch {epoch}  loss {net.loss(y_onehot):.4f}")

    preds = np.argmax(net.forward(X), axis=1)
    print("Predictions:", preds, " True:", y)
