class SimpleNeuralNetwork:
    """A simple 2-layer neural network implementation"""

    def __init__(self, input_size=2, hidden_size=4, output_size=1):
        # Initialize weights
        self.W1 = np.random.randn(input_size, hidden_size) * 0.5
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.5
        self.b2 = np.zeros((1, output_size))

    def sigmoid(self, x):
        return 1 / (1 + np.exp(-np.clip(x, -500, 500)))

    def forward(self, X):
        self.z1 = X @ self.W1 + self.b1
        self.a1 = self.sigmoid(self.z1)
        self.z2 = self.a1 @ self.W2 + self.b2
        self.a2 = self.sigmoid(self.z2)
        return self.a2

    def backward(self, X, y):
        m = X.shape[0]

        # Output layer gradients
        dz2 = 2 * (self.a2 - y) * self.sigmoid(self.z2) * (1 - self.sigmoid(self.z2))
        dW2 = self.a1.T @ dz2 / m
        db2 = np.mean(dz2, axis=0, keepdims=True)

        # Hidden layer gradients
        dz1 = (dz2 @ self.W2.T) * self.sigmoid(self.z1) * (1 - self.sigmoid(self.z1))
        dW1 = X.T @ dz1 / m
        db1 = np.mean(dz1, axis=0, keepdims=True)

        return dW1, db1, dW2, db2

    def train(self, X, y, epochs=1000, learning_rate=1.0):
        losses = []

        for epoch in range(epochs):
            # Forward pass
            predictions = self.forward(X)

            # Compute loss
            loss = np.mean((predictions - y) ** 2)
            losses.append(loss)

            # Backward pass
            dW1, db1, dW2, db2 = self.backward(X, y)

            # Update weights
            self.W1 -= learning_rate * dW1
            self.b1 -= learning_rate * db1
            self.W2 -= learning_rate * dW2
            self.b2 -= learning_rate * db2

            if epoch % 100 == 0:
                print(f"Epoch {epoch:4d}: Loss = {loss:.6f}")

        return losses

    def predict(self, X):
        return self.forward(X)

# Test the class
nn = SimpleNeuralNetwork()
losses = nn.train(X, y, epochs=1000, learning_rate=1.0)
predictions = nn.predict(X)

print("\nClass-based Neural Network Results:")
for i in range(len(X)):
    pred = predictions[i, 0]
    target = y[i, 0]
    print(f"{X[i]} -> {target} | Prediction: {pred:.4f} | Rounded: {round(pred)}")