import numpy as np

META = {'name': 'kpz_field_forecast', 'domain': 'pde', 'description': 'One-step forecasting of a periodic 3-D KPZ field.'}

def _step(h, nu=0.22, lam=1.8, dt=0.025):
    lap = np.zeros_like(h)
    grad2 = np.zeros_like(h)
    for ax in (2, 3, 4):
        hp, hm = np.roll(h, -1, ax), np.roll(h, 1, ax)
        lap += hp - 2*h + hm
        grad2 += ((hp-hm)/2.0)**2
    return h + dt*(nu*lap + 0.5*lam*grad2)

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    n = n_train + n_test
    h = rng.normal(size=(n,1,8,8,8)).astype(np.float32)
    for ax in (2,3,4):
        h = 0.55*h + 0.15*(np.roll(h,1,ax)+np.roll(h,-1,ax))
    h /= h.std(axis=(2,3,4), keepdims=True) + 1e-6
    y = _step(h).astype(np.float32)
    return {'xtr':h[:n_train], 'ytr':y[:n_train], 'xte':h[n_train:], 'yte':y[n_train:], 'task':'regression', 'metric':'mse', 'out_dim':1}