import numpy as np

META = {"name": "poisson_measurements", "domain": "pde", "description": "Synthetic Poisson-like elliptic field operator with irregular sensor/query meshes."}


def _fields(seed, n, m=16):
    rng = np.random.RandomState(seed)
    # Smooth forcing fields and a stable spectral inverse (Poisson surrogate).
    xx, yy = np.meshgrid(np.arange(m) / m, np.arange(m) / m, indexing="ij")
    outx, outy = [], []
    for _ in range(n):
        f = np.zeros((m, m), np.float32)
        for k in range(1, 4):
            for l in range(1, 4):
                a = rng.randn() / (k*k + l*l)
                f += a * np.sin(2*np.pi*k*xx) * np.sin(2*np.pi*l*yy)
        # Smooth elliptic solution, exactly generated in a low-frequency basis.
        u = np.zeros_like(f)
        for k in range(1, 4):
            for l in range(1, 4):
                # recover coefficients by projection (basis is orthogonal on grid)
                basis = np.sin(2*np.pi*k*xx) * np.sin(2*np.pi*l*yy)
                c = float((f*basis).mean()) / max(float((basis*basis).mean()), 1e-8)
                u += c / (k*k+l*l) * basis
        outx.append(f.reshape(-1)); outy.append(u.reshape(-1))
    return np.asarray(outx, np.float32), np.asarray(outy, np.float32)


def get_dataset(seed, n_train, n_test):
    xtr, ytr = _fields(seed + 17, n_train)
    xte, yte = _fields(seed + 991, n_test)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "input_shape": (256,), "out_dim": 256}