import numpy as np
META = {"name":"symmetric_periodic_poisson","domain":"pde","description":"2-D periodic Poisson inverse on an 8x8 square grid; source fields and exact zero-mean solutions."}

def _lap_eigs(n):
    k=np.arange(n)
    return (4-2*np.cos(2*np.pi*k/n)[:,None]-2*np.cos(2*np.pi*k/n)[None,:]).astype(np.float32)

def get_dataset(seed, n_train, n_test):
    def sample(n, s):
        rng=np.random.RandomState(s); N=8
        # Random Fourier sources with no constant mode, ensuring solvability.
        f=rng.normal(size=(n,N,N)).astype(np.float32)
        fh=np.fft.fft2(f,axes=(-2,-1)); fh[:,0,0]=0
        lam=_lap_eigs(N); uh=fh/(lam[None]+1e-5); uh[:,0,0]=0
        u=np.fft.ifft2(uh,axes=(-2,-1)).real.astype(np.float32)
        # Normalize each sample to avoid scale-driven triviality.
        f=f/(f.std(axis=(1,2),keepdims=True)+1e-4); u=u/(u.std(axis=(1,2),keepdims=True)+1e-4)
        return f.reshape(n,-1),u.reshape(n,-1)
    xtr,ytr=sample(n_train,seed); xte,yte=sample(n_test,seed+5000)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":64}