import numpy as np

META = {
    "name": "simplicial_cochain_regression",
    "domain": "simplicial_geometry_pde_like",
    "description": "Regression from noisy node and edge cochains on a filled oriented triangle; targets are compatible cochain observables."
}

D0 = np.array([[-1., 1., 0.], [-1., 0., 1.], [0., -1., 1.]])
D1 = np.array([[1., -1., 1.]])


def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.RandomState(s)
        v = rng.normal(size=(n, 3, 2)).astype(np.float32)
        e = np.einsum('ev,nvc->nec', D0, v).astype(np.float32)
        x0 = v + 0.65 * rng.normal(size=v.shape).astype(np.float32)
        x1 = e + 0.65 * rng.normal(size=e.shape).astype(np.float32)
        x = np.concatenate([x0, x1], axis=1)
        y = (v[:, :, 0].mean(axis=1) + 0.35 * e[:, :, 1].mean(axis=1)
             + 0.15 * np.tanh(v[:, :, 1].sum(axis=1))).astype(np.float32)
        return x, y[:, None]
    xtr, ytr = make(n_train, seed)
    xte, yte = make(n_test, seed + 5000)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}