import numpy as np

META = {
    "name": "rough_energy_regression",
    "domain": "optimization",
    "description": "Regression on a smooth baseline plus a finite Weierstrass rough component; used to test dilation-matched updates in a neural training system."
}

A, B, N, K = 0.8, 3.0, 12, 0.25

def phi(x):
    return np.sin(2*np.pi*x)

def rough(x):
    x = np.asarray(x, dtype=np.float64)
    z = np.zeros(x.shape[0])
    for n in range(N+1):
        z += A**n * phi(B**n*x[:, 0])
    return z

def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(int(seed))
    xtr = rng.uniform(0.2, 3.8, size=(n_train, 4)).astype(np.float32)
    xte = rng.uniform(0.2, 3.8, size=(n_test, 4)).astype(np.float32)
    def target(x):
        smooth = 0.35*(x[:,0]-2.0)**2 + 0.15*x[:,1] - 0.10*x[:,2] + 0.08*x[:,3]
        return (smooth + K*rough(x)).astype(np.float32)
    return {"xtr":xtr, "ytr":target(xtr)[:,None], "xte":xte,
            "yte":target(xte)[:,None], "task":"regression", "metric":"mse",
            "input_shape":(4,), "out_dim":1}