# Create a grid of points to visualize the decision boundary
def plot_decision_boundary(W1, b1, W2, b2):
    """Plot the decision boundary learned by the network"""
    # Create a grid
    xx, yy = np.meshgrid(np.linspace(-0.5, 1.5, 100),
                         np.linspace(-0.5, 1.5, 100))

    # Flatten the grid for prediction
    grid_points = np.c_[xx.ravel(), yy.ravel()]

    # Make predictions on the grid
    _, _, _, grid_predictions = forward_pass(grid_points, W1, b1, W2, b2)
    grid_predictions = grid_predictions.reshape(xx.shape)

    # Plot
    plt.figure(figsize=(8, 6))
    plt.contourf(xx, yy, grid_predictions, levels=50, alpha=0.8, cmap='RdYlBu')
    plt.colorbar(label='Network Output')

    # Plot data points
    colors = ['red' if label == 0 else 'blue' for label in y.flatten()]
    plt.scatter(X[:, 0], X[:, 1], c=colors, s=100, edgecolors='black', linewidth=2)

    # Add labels
    for i, (x, y_val) in enumerate(zip(X, y.flatten())):
        plt.annotate(f'({x[0]},{x[1]})→{y_val}', 
                    (x[0], x[1]), xytext=(5, 5), textcoords='offset points')

    plt.title('Neural Network Decision Boundary')
    plt.xlabel('Input 1')
    plt.ylabel('Input 2')
    plt.grid(True, alpha=0.3)
    plt.show()

# Visualize the decision boundary
plot_decision_boundary(W1_trained, b1_trained, W2_trained, b2_trained)