import numpy as np


def normalize_adjacency(A):
    A_tilde = A + np.eye(A.shape[0])
    degrees = A_tilde.sum(axis=1)
    D_inv_sqrt = np.diag(1.0 / np.sqrt(degrees))
    return D_inv_sqrt @ A_tilde @ D_inv_sqrt


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


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


class GCNLayer:
    def __init__(self, in_dim, out_dim, rng):
        scale = np.sqrt(2.0 / in_dim)
        self.W = rng.normal(0, scale, size=(in_dim, out_dim))

    def forward(self, A_norm, H):
        self.H_in = H
        self.A_norm = A_norm
        Z = A_norm @ H @ self.W
        return Z

    def backward(self, dZ, lr):
        # dZ: gradient wrt this layer's pre-activation output
        AH = self.A_norm @ self.H_in
        grad_W = AH.T @ dZ
        dH = self.A_norm.T @ (dZ @ self.W.T)
        self.W -= lr * grad_W
        return dH


class TwoLayerGCN:
    def __init__(self, in_dim, hidden_dim, n_classes, seed=0):
        rng = np.random.default_rng(seed)
        self.layer1 = GCNLayer(in_dim, hidden_dim, rng)
        self.layer2 = GCNLayer(hidden_dim, n_classes, rng)

    def forward(self, A_norm, X):
        self.z1 = self.layer1.forward(A_norm, X)
        self.h1 = relu(self.z1)
        self.z2 = self.layer2.forward(A_norm, self.h1)
        self.probs = softmax(self.z2)
        return self.probs

    def backward(self, y_onehot, lr):
        n = y_onehot.shape[0]
        dZ2 = (self.probs - y_onehot) / n
        dH1 = self.layer2.backward(dZ2, lr)
        dZ1 = dH1 * (self.z1 > 0)
        self.layer1.backward(dZ1, lr)

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