import numpy as np


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


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


def softmax(z):
    z = z - np.max(z, axis=1, keepdims=True)
    exp_z = np.exp(z)
    return exp_z / exp_z.sum(axis=1, keepdims=True)


class TwoLayerNet:
    def __init__(self, d_in, d_hidden, d_out, seed=0):
        rng = np.random.default_rng(seed)
        scale1 = np.sqrt(2.0 / d_in)
        scale2 = np.sqrt(2.0 / d_hidden)
        self.W1 = rng.normal(0, scale1, size=(d_in, d_hidden))
        self.b1 = np.zeros(d_hidden)
        self.W2 = rng.normal(0, scale2, size=(d_hidden, d_out))
        self.b2 = np.zeros(d_out)

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

    def backward(self, X, y_onehot, lr=0.1):
        n = X.shape[0]
        delta2 = self.a2 - y_onehot  # softmax + cross-entropy gradient
        grad_W2 = self.a1.T @ delta2 / n
        grad_b2 = delta2.mean(axis=0)

        delta1 = (delta2 @ self.W2.T) * relu_grad(self.z1)
        grad_W1 = X.T @ delta1 / n
        grad_b1 = delta1.mean(axis=0)

        self.W2 -= lr * grad_W2
        self.b2 -= lr * grad_b2
        self.W1 -= lr * grad_W1
        self.b1 -= lr * grad_b1

    def loss(self, y_onehot):
        eps = 1e-12
        return -np.mean(np.sum(y_onehot * np.log(self.a2 + eps), axis=1))
