import numpy as np
META = {"name":"transport_density","domain":"transport","description":"1D nonnegative transported densities: reconstruct the next state from a density input."}

def _density(x, c, w):
    y=np.exp(-0.5*((x-c)/w)**2)+0.35*np.exp(-0.5*((x-(c-.14))/(.65*w))**2)
    y=np.maximum(y,1e-12); y/=np.trapz(y,x); return y

def get_dataset(seed, n_train, n_test):
    rng=np.random.RandomState(seed); x=np.linspace(0,1,32,dtype=np.float32)
    def make(n):
        U=[]; V=[]
        for _ in range(n):
            c=rng.uniform(.18,.82); w=rng.uniform(.035,.09); shift=rng.uniform(-.10,.10)
            u=_density(x,c,w); v=_density(x,np.clip(c+shift,.08,.92),w)
            U.append(u.astype(np.float32)); V.append(v.astype(np.float32))
        return np.asarray(U),np.asarray(V)
    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":32}