import numpy as np

META = {
    'name': 'quadratic_context_multitoken',
    'domain': 'diffusion-sampling',
    'description': 'Multitoken denoising with transferable analytic quadratic Gaussian context.'
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    L = 8
    def make(n):
        t = np.linspace(0, 1, L, dtype=np.float32)
        z = rng.choice([-1., 1.], size=n).astype(np.float32)
        phase = rng.uniform(-np.pi, np.pi, n).astype(np.float32)
        amp = rng.uniform(.8, 1.2, n).astype(np.float32)
        x = np.empty((n, L), dtype=np.float32)
        for i in range(n):
            x[i] = amp[i] * (np.sin(2*np.pi*1.25*t + phase[i]) + .35*z[i]*np.cos(2*np.pi*2.0*t))
        return x
    rtr, rte = make(n_train), make(n_test)
    k = 1.5
    return {'xtr': rtr, 'ytr': rtr/(1+k), 'xte': rte, 'yte': rte/(1+k),
            'task': 'regression', 'metric': 'mse', 'out_dim': L, 'k': k}