import numpy as np

META = {
    'name': 'fisher_multitoken_denoising',
    'domain': 'diffusion-sampling',
    'description': 'Correlated multi-token Gaussian sequences for neural denoising under an annealed noise path.'
}

def get_dataset(seed, n_train, n_test):
    L = 12
    def make(n, rs):
        t = np.linspace(0.0, 1.0, L, dtype=np.float32)
        out = np.empty((n, L), dtype=np.float32)
        for i in range(n):
            a = rs.uniform(0.6, 1.4)
            f = rs.uniform(0.7, 1.5)
            ph = rs.uniform(0.0, 2.0*np.pi)
            z = a*np.sin(2.0*np.pi*f*t + ph) + 0.35*np.cos(np.pi*f*t - 0.4*ph)
            out[i] = z + rs.normal(0.0, 0.06, L)
        return out.astype(np.float32)
    xtr = make(n_train, np.random.RandomState(seed))
    xte = make(n_test, np.random.RandomState(seed + 10007))
    return {'xtr': xtr, 'ytr': xtr.copy(), 'xte': xte, 'yte': xte.copy(),
            'task': 'regression', 'metric': 'mse', 'out_dim': L}