import numpy as np

META = {
    "name": "separable_diffusion_field",
    "domain": "pde",
    "description": "3D voxel fields from an anisotropic separable diffusion map, with scalar energy target."
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    shape = (4, 4, 4)
    q = np.arange(4, dtype=np.float32)
    z = np.exp(-0.5 * ((q[:, None] - q[None, :]) / 0.9) ** 2)
    z /= z.sum(axis=1, keepdims=True)
    w = np.linspace(-1.0, 1.0, 64).astype(np.float32)

    def make(n):
        x = rng.normal(0, 1, (n,) + shape).astype(np.float32)
        y = np.einsum('ab,nbcd->nacd', z, x)
        y = np.einsum('ab,nacd->nbcd', z, y)
        y = np.einsum('ab,nbcd->nabc', z, y)
        # Scalar PDE observable; nonlinear term prevents a trivial identity rule.
        target = (y.reshape(n, 64) @ w / 8.0 + 0.15 * np.mean(np.tanh(x), axis=(1, 2, 3)))
        return x, target.astype(np.float32)[:, None]

    xtr, ytr = make(n_train)
    xte, yte = make(n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 1}