import numpy as np

META = {
    'name': 'bloch_gauge_pde',
    'domain': 'pde',
    'description': '1D periodic Bloch-factor field surrogate with exact reciprocal-lattice gauge covariance.'
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(int(seed))
    n = int(n_train) + int(n_test)
    c = rng.uniform(-1.0, 1.0, n).astype('float32')
    x = rng.uniform(0.0, 1.0, n).astype('float32')
    q0 = rng.uniform(-np.pi, np.pi, n).astype('float32')
    m = rng.integers(-2, 3, n).astype('float32')
    q = q0 + m * (2*np.pi)
    amp = np.exp(0.30j*c*np.cos(2*np.pi*x) + 0.10j*c*np.sin(4*np.pi*x))
    f0 = amp * (1 + 0.12*q0*np.cos(2*np.pi*x) + 0.05j*q0*np.sin(2*np.pi*x))
    y = f0 * np.exp(-1j*m*2*np.pi*x)
    X = np.stack([c, x, q], axis=1).astype('float32')
    Y = np.stack([y.real, y.imag], axis=1).astype('float32')
    return {'xtr': X[:n_train], 'ytr': Y[:n_train], 'xte': X[n_train:], 'yte': Y[n_train:], 'task': 'regression', 'metric': 'mse', 'out_dim': 2}