def make_two_community_graph(n_per_community=15, seed=0):
    rng = np.random.default_rng(seed)
    n = 2 * n_per_community
    A = np.zeros((n, n))

    # dense connections within each community, sparse across
    for i in range(n):
        for j in range(i + 1, n):
            same_community = (i < n_per_community) == (j < n_per_community)
            p = 0.3 if same_community else 0.02
            if rng.random() < p:
                A[i, j] = A[j, i] = 1.0

    labels = np.array(
        [0] * n_per_community + [1] * n_per_community
    )
    # node features: mostly noise, so the graph structure -- not
    # the raw features -- has to carry most of the signal
    X = rng.normal(0, 1, size=(n, 8))
    return A, X, labels


def one_hot(y, n_classes):
    out = np.zeros((len(y), n_classes))
    out[np.arange(len(y)), y] = 1
    return out


if __name__ == "__main__":
    A, X, labels = make_two_community_graph()
    A_norm = normalize_adjacency(A)
    y_onehot = one_hot(labels, n_classes=2)

    model = TwoLayerGCN(in_dim=8, hidden_dim=16, n_classes=2)
    for epoch in range(200):
        probs = model.forward(A_norm, X)
        loss = model.loss(y_onehot)
        model.backward(y_onehot, lr=0.5)
        if epoch % 50 == 0:
            preds = probs.argmax(axis=1)
            acc = (preds == labels).mean()
            print(f"epoch {epoch}  loss {loss:.4f}  acc {acc:.3f}")

    final_preds = model.forward(A_norm, X).argmax(axis=1)
    final_acc = (final_preds == labels).mean()
    print(f"Final training accuracy: {final_acc:.3f}")
