import numpy as np

META = {
    'name': 'conditional_sequence_diffusion',
    'domain': 'sequence-level diffusion sampling',
    'description': 'Conditional generation of multi-token sinusoidal sequences from a context window; the metric is held-out conditional reconstruction MSE.'
}

def get_dataset(seed, n_train=400, n_test=100):
    rng = np.random.default_rng(int(seed))
    n = n_train + n_test
    ctx = rng.uniform(-1, 1, (n, 4)).astype('float32')
    phase = rng.uniform(0, 2*np.pi, n)
    freq = rng.uniform(.7, 1.8, n)
    amp = .7 + .5 * (ctx[:, 0] + 1) / 2
    tt = np.arange(8, dtype='float32')[None, :]
    y = (amp[:, None] * np.sin(freq[:, None] * tt + phase[:, None]) +
         .25 * ctx[:, 1, None] * np.cos(.55 * tt + ctx[:, 2, None]) +
         .08 * rng.normal(size=(n, 8))).astype('float32')
    # Context contains informative phase/frequency proxies, making this a conditional task.
    ctx[:, 2] = np.sin(phase); ctx[:, 3] = np.cos(phase)
    return {'xtr': ctx[:n_train], 'ytr': y[:n_train],
            'xte': ctx[n_train:], 'yte': y[n_train:], 'task': 'regression', 'metric': 'mse',
            'input_shape': (4,), 'out_dim': 8, 'track': META['name']}