import math
import numpy as np

META = {'name':'two_view_gauge_localization_v2','domain':'geometry/localization','description':'Planar two-view global and relay-frame observations for unknown yaw and translation, predicting target x-coordinate.'}

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])
    for z in (q1,q2,ev1,ev2,et): z += rng.normal(0,.035,z.shape).astype(np.float32)
    return np.concatenate([q1,q2,ev1,ev2,et],1).astype(np.float32),p[:,0:1].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':1}