import numpy as np

META = {
    "name": "poisson_field_operator",
    "domain": "pde",
    "description": "Small manufactured 2-D Poisson/BVP operator: Fourier forcing coefficients to a spatial solution field.",
}


def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    side = 8
    xx, yy = np.meshgrid(np.linspace(0, 1, side), np.linspace(0, 1, side), indexing="ij")
    modes = [(1, 1), (1, 2), (2, 1), (2, 2), (1, 3), (3, 1)]

    def sample(n, rs):
        z = rs.normal(size=(n, len(modes))).astype(np.float32)
        fields = []
        for i, (kx, ky) in enumerate(modes):
            basis = np.sin(np.pi * kx * xx) * np.sin(np.pi * ky * yy)
            # Inverse-Laplacian scaling is the manufactured Poisson structure.
            fields.append(z[:, i, None] * basis.reshape(1, -1) / (kx * kx + ky * ky))
        u = np.sum(fields, axis=0).astype(np.float32)
        # Heteroscedastic observation/model discrepancy is localized near a boundary.
        noise_scale = (0.012 + 0.030 * np.exp(-((xx - .15) ** 2 + (yy - .15) ** 2) / .025)).reshape(1, -1)
        u += rs.normal(size=u.shape).astype(np.float32) * noise_scale
        return z, u

    xtr, ytr = sample(n_train, rng)
    xte, yte = sample(n_test, np.random.RandomState(seed + 5000))
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": side * side}