import numpy as np

META = {
    'name': 'directed_path_holonomy',
    'domain': 'graph',
    'description': 'Directed length-3 path composition regression with noncommutative edge operators; order is essential.'
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))

    def make(n):
        x = rng.normal(0.0, 0.45, size=(n, 3, 2, 2)).astype(np.float32)
        x += np.eye(2, dtype=np.float32)[None, None, :, :] * 0.8
        h = np.matmul(np.matmul(x[:, 0], x[:, 1]), x[:, 2])
        y = h.reshape(n, 4).astype(np.float32)
        return x.reshape(n, 12), 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': (12,), 'out_dim': 4,
    }