"""
Lecture 11 - Building Neural Networks in Python: NumPy from scratch

Generalizes Lecture 8's hand-derived, single-example backprop to a small
batch, using the exact vectorized update equations from Lecture 8's
"Vectorized Form (Mini-Batches)" section. Architecture: 2 inputs -> 2
hidden units (ReLU) -> 1 output (sigmoid), binary cross-entropy loss --
identical to the architecture used throughout Lectures 8-9.

Trained on the AND-gate toy dataset (the same points introduced in
Lecture 3).

Run: python lecture-11-numpy-nn.py
"""
import numpy as np


def sigmoid(z):
    return 1 / (1 + np.exp(-z))


def relu(z):
    return np.maximum(0, z)


def relu_grad(z):
    return (z > 0).astype(float)


class NeuralNetwork:
    """2 inputs -> 2 hidden (ReLU) -> 1 output (sigmoid), same architecture as Lecture 8."""

    def __init__(self, seed=42):
        rng = np.random.default_rng(seed)
        self.W1 = rng.normal(0, 0.5, size=(2, 2))
        self.b1 = np.zeros((1, 2))
        self.W2 = rng.normal(0, 0.5, size=(2, 1))
        self.b2 = np.zeros((1, 1))

    def forward(self, X):
        self.X = X
        self.Z1 = X @ self.W1 + self.b1
        self.A = relu(self.Z1)
        self.Z2 = self.A @ self.W2 + self.b2
        self.Q = sigmoid(self.Z2)
        return self.Q

    def backward(self, y, lr):
        m = y.shape[0]
        dZ2 = self.Q - y                           # sigmoid + BCE shortcut (Lecture 8)
        dW2 = self.A.T @ dZ2 / m
        db2 = np.sum(dZ2, axis=0, keepdims=True) / m
        dA = dZ2 @ self.W2.T
        dZ1 = dA * relu_grad(self.Z1)
        dW1 = self.X.T @ dZ1 / m
        db1 = np.sum(dZ1, axis=0, keepdims=True) / m
        self.W1 -= lr * dW1; self.b1 -= lr * db1   # gradient descent (Lecture 7)
        self.W2 -= lr * dW2; self.b2 -= lr * db2

    def train(self, X, y, lr=0.5, epochs=1000, verbose=True):
        for epoch in range(epochs):
            Q = self.forward(X)
            self.backward(y, lr)
            if verbose and epoch % 100 == 0:
                loss = -np.mean(y * np.log(Q + 1e-9) + (1 - y) * np.log(1 - Q + 1e-9))
                print(f"epoch {epoch:4d}  loss={loss:.4f}")
        return self


if __name__ == "__main__":
    # AND-gate toy dataset (Lecture 3's points)
    X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
    y = np.array([[0], [0], [0], [1]], dtype=float)

    nn = NeuralNetwork(seed=42)
    nn.train(X, y, lr=0.5, epochs=1000)

    print("\nFinal predictions:", nn.forward(X).ravel().round(4))
    print("True labels:      ", y.ravel())