import math
import numpy as np

META = {
    'name': 'two_view_gauge_localization',
    'domain': 'geometry/localization',
    'description': 'Planar two-view global and relay-frame observations for target localization with an unknown yaw and translation gauge.'
}

def _rot(a):
    c, s = np.cos(a), np.sin(a)
    return np.array([[c, -s], [s, c]], dtype=np.float32)

def _make(seed, n):
    rng = np.random.default_rng(seed)
    x = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
    p = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
    psi = rng.uniform(-math.pi, math.pi, n)
    q1 = rng.uniform(-1, 1, (n, 2)).astype(np.float32)
    dq = rng.uniform(-2, 2, (n, 2)).astype(np.float32)
    dq += (np.linalg.norm(dq, axis=1) < .25)[:, None] * np.array([.4, .1], np.float32)
    q2 = q1 + dq
    ev1 = np.empty_like(q1); ev2 = np.empty_like(q1); et = np.empty_like(q1)
    for i, a in enumerate(psi):
        R = _rot(-a)
        ev1[i] = R @ (q1[i] - x[i]); ev2[i] = R @ (q2[i] - x[i]); et[i] = R @ (p[i] - x[i])
    noise = .035
    for z in (q1, q2, ev1, ev2, et):
        z += rng.normal(0, noise, z.shape).astype(np.float32)
    return np.concatenate([q1, q2, ev1, ev2, et], axis=1).astype(np.float32), p.astype(np.float32)

def get_dataset(seed, n_train, n_test):
    xtr, ytr = _make(seed, n_train)
    xte, yte = _make(seed + 5000, n_test)
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
            'task': 'regression', 'metric': 'mse', 'out_dim': 2}