import numpy as np

META = {
    "name": "geometric_graph_diffusion",
    "domain": "graph-neural-networks",
    "description": "Node-field regression on geometric kNN graphs; target is a short distance-weighted diffusion response."
}


def _one(seed, n):
    rng = np.random.RandomState(seed)
    xy = rng.uniform(-1.0, 1.0, size=(24, 2)).astype(np.float32)
    d2 = ((xy[:, None] - xy[None, :]) ** 2).sum(-1)
    adj = np.zeros_like(d2, dtype=np.float32)
    for i in range(24):
        adj[i, np.argsort(d2[i])[1:6]] = 1.0
    adj = np.maximum(adj, adj.T)
    w = adj * np.maximum(d2, 1e-5)  # n=4: d^(n-2) = d^2
    p = w / np.maximum(w.sum(1, keepdims=True), 1e-8)
    q = 0.55 * np.eye(24, dtype=np.float32) + 0.45 * p
    kernel = q @ q
    root = int(np.argmin((xy ** 2).sum(1)))
    x = rng.normal(size=(n, 24, 2)).astype(np.float32)
    y = (x[:, :, 0] @ kernel[root]).astype(np.float32)
    y += 0.08 * rng.normal(size=n).astype(np.float32)
    coords = np.broadcast_to(xy, (n, 24, 2)).copy()
    root_flag = np.zeros((n, 24, 1), dtype=np.float32)
    root_flag[:, root, 0] = 1.0
    return np.concatenate([x, coords, root_flag], axis=2)


def get_dataset(seed, n_train=400, n_test=160):
    # The graph is regenerated deterministically per split; coordinates and
    # root marker are part of every sample, so the task is self-contained.
    # Targets are generated with the same weighted diffusion rule as the idea,
    # while evaluation remains ordinary regression MSE.
    def make(seed0, n):
        rng = np.random.RandomState(seed0)
        xy = rng.uniform(-1.0, 1.0, size=(24, 2)).astype(np.float32)
        d2 = ((xy[:, None] - xy[None, :]) ** 2).sum(-1)
        adj = np.zeros_like(d2, dtype=np.float32)
        for i in range(24):
            adj[i, np.argsort(d2[i])[1:6]] = 1.0
        adj = np.maximum(adj, adj.T)
        w = adj * np.maximum(d2, 1e-5)
        p = w / np.maximum(w.sum(1, keepdims=True), 1e-8)
        q = 0.55 * np.eye(24, dtype=np.float32) + 0.45 * p
        kernel = q @ q
        root = int(np.argmin((xy ** 2).sum(1)))
        x = rng.normal(size=(n, 24, 2)).astype(np.float32)
        y = (x[:, :, 0] @ kernel[root]).astype(np.float32)
        y += 0.08 * rng.normal(size=n).astype(np.float32)
        coords = np.broadcast_to(xy, (n, 24, 2)).copy()
        flag = np.zeros((n, 24, 1), dtype=np.float32); flag[:, root, 0] = 1.0
        return np.concatenate([x, coords, flag], axis=2), y
    xtr, ytr = make(seed, n_train)
    xte, yte = make(seed + 5000, n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}