"""Local custom track: supervised harmonic Dirichlet/Poisson-style collocation."""
import numpy as np

META = {
    "name": "poisson_harmonic_square",
    "domain": "pde",
    "description": "Unit-square collocation for the harmonic Dirichlet solution u(x,y)=x^2-y^2; interior PDE residual is Delta u=0 and values provide boundary/data supervision."
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    # Keep the train/test distributions identical while retaining a genuine
    # collocation domain; the PDE loss is evaluated on every training point.
    xtr = rng.uniform(-1.0, 1.0, size=(int(n_train), 2)).astype(np.float32)
    xte = rng.uniform(-1.0, 1.0, size=(int(n_test), 2)).astype(np.float32)
    def solution(x):
        return (x[:, 0] ** 2 - x[:, 1] ** 2)[:, None].astype(np.float32)
    return {"xtr": xtr, "ytr": solution(xtr), "xte": xte, "yte": solution(xte),
            "task": "regression", "metric": "mse", "input_shape": (2,), "out_dim": 1}