import numpy as np

META = {'name':'tangent_surface_vector_regression','domain':'geometry/point-cloud','description':'Noisy rotated sphere point clouds with a tangent-field magnitude target; the input contains ambient vector features with normal corruption.'}

def _rot(rng):
    q,r=np.linalg.qr(rng.normal(size=(3,3)))
    q=q@np.diag(np.sign(np.diag(r)))
    if np.linalg.det(q)<0:q[:,0]*=-1
    return q.astype(np.float32)

def _make(seed,n):
    rng=np.random.default_rng(seed); m=20; X=np.empty((n,m,7),np.float32); Y=np.empty((n,1),np.float32)
    a=np.array([.8,-.35,.55],np.float32)
    for b in range(n):
        z=rng.normal(size=(m,3)).astype(np.float32); z/=np.linalg.norm(z,axis=1,keepdims=True)
        R=_rot(rng); p=z@R.T; p[1:]=p[0]+.55*(z[1:]-z[0])@R.T; p[1:]/=np.linalg.norm(p[1:],axis=1,keepdims=True)
        s=p@a; v=a[None,:]-s[:,None]*p
        v=v+.9*rng.normal(size=(m,1)).astype(np.float32)*p
        X[b,:,:3]=p; X[b,:,3]=s; X[b,:,4:]=v; Y[b,0]=np.linalg.norm(v[0]-((v[0]@p[0])*p[0]))
    return X,Y

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}