import numpy as np

META = {
    "name": "graph_landmark_forecast",
    "domain": "graph-nn",
    "description": "Directed graph node forecasting: predict one-step neighborhood propagation from node signals; landmark signatures are valid structural positional features."
}

# Fixed directed graph, deliberately with distinct incoming signatures for landmarks.
def _graph(n=20):
    rng = np.random.default_rng(731)
    A = (rng.random((n, n)) < 0.22).astype(np.float32)
    np.fill_diagonal(A, 0)
    # Avoid empty incoming neighborhoods and make the graph nondegenerate.
    for v in range(n):
        if A[:, v].sum() == 0:
            A[rng.integers(n), v] = 1
    return A


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed) + 1200)
    A = _graph()
    n = A.shape[0]
    deg = A.sum(0, keepdims=True).T
    P = A / np.maximum(deg, 1.0)
    total = n_train + n_test
    # Inputs are independent node signals; targets require the directed graph.
    x = rng.normal(size=(total, n, 1)).astype(np.float32)
    y = (0.55*x[:, :, 0] + 0.45*np.einsum('vu,bu->bv', P, x[:, :, 0])).astype(np.float32)
    return {"xtr": x[:n_train], "ytr": y[:n_train, :, None],
            "xte": x[n_train:], "yte": y[n_train:, :, None],
            "task": "regression", "metric": "mse", "out_dim": n,
            "adjacency": A}