import numpy as np

META = {
    'name': 'mixed_junction_pde',
    'domain': 'pde',
    'description': 'Manufactured mixed Dirichlet-Neumann junction field on a square; the corner at the origin has Kondratjev r^1/2 sin(theta/2) regularity.'
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    # Positive quadrant is a local model of a straight mixed junction.
    xtr = rng.uniform(0.0, 1.0, (n_train, 2)).astype(np.float32)
    xte = rng.uniform(0.0, 1.0, (n_test, 2)).astype(np.float32)
    def field(x):
        r = np.sqrt(np.sum(x*x, axis=1) + 1e-14)
        th = np.arctan2(x[:, 1], x[:, 0])
        # Singular harmonic mode for a mixed corner, plus smooth harmonic-like terms.
        return (r**0.5 * np.sin(0.5*th) +
                0.20*np.sin(np.pi*x[:, 0])*np.sin(np.pi*x[:, 1]) +
                0.10*x[:, 0]*x[:, 1]).astype(np.float32)
    return {'xtr': xtr, 'ytr': field(xtr).reshape(-1, 1), 'xte': xte, 'yte': field(xte).reshape(-1, 1),
            'task': 'regression', 'metric': 'mse', 'out_dim': 1}