# Test different hidden layer sizes
hidden_sizes = [2, 4, 8, 16]
results = {}

for hidden_size in hidden_sizes:
    print(f"\nTesting hidden size: {hidden_size}")
    nn = SimpleNeuralNetwork(input_size=2, hidden_size=hidden_size, output_size=1)
    losses = nn.train(X, y, epochs=1000, learning_rate=1.0)
    predictions = nn.predict(X)

    # Calculate accuracy
    accuracy = np.mean(np.round(predictions) == y)
    results[hidden_size] = {
        'final_loss': losses[-1],
        'accuracy': accuracy,
        'predictions': predictions
    }

    print(f"  Final loss: {losses[-1]:.6f}")
    print(f"  Accuracy: {accuracy:.1%}")

# Plot comparison
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
hidden_sizes_list = list(results.keys())
final_losses = [results[hs]['final_loss'] for hs in hidden_sizes_list]
plt.bar(hidden_sizes_list, final_losses)
plt.title('Final Loss vs Hidden Size')
plt.xlabel('Hidden Layer Size')
plt.ylabel('Final Loss')
plt.yscale('log')

plt.subplot(1, 2, 2)
accuracies = [results[hs]['accuracy'] for hs in hidden_sizes_list]
plt.bar(hidden_sizes_list, accuracies)
plt.title('Accuracy vs Hidden Size')
plt.xlabel('Hidden Layer Size')
plt.ylabel('Accuracy')
plt.ylim(0, 1)

plt.tight_layout()
plt.show()