"""Custom benchmark track: intermittent stable pair relations in dynamics."""
META = {"name":"relational_dynamics", "domain":"dynamics", "description":"Multi-agent intermittent-contact regression with stable pair identities and delayed reactivation."}
import numpy as np
T, P, F = 12, 6, 3

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    def make(n):
        latent = rng.normal(0, 1, (n, P)).astype(np.float32)
        x = rng.normal(0, .08, (n, T, P, 2)).astype(np.float32)
        active = np.zeros((n, T, P), dtype=np.float32)
        active[:, 0:3, :] = 1
        active[:, 9:12, :] = 1
        x[:, 0:3, :, 0] += latent[:, None, :]
        x[:, 0:3, :, 1] += .25 * latent[:, None, :]
        seq = np.concatenate([x, active[..., None]], axis=-1)
        y = (latent.sum(axis=1, keepdims=True) / np.sqrt(P)).astype(np.float32)
        return seq.reshape(n, -1), y
    xtr, ytr = make(n_train); xte, yte = make(n_test)
    return {"xtr":xtr, "ytr":ytr, "xte":xte, "yte":yte,
            "task":"regression", "metric":"mse", "input_shape":(T*P*F,), "out_dim":1}