import numpy as np

META = {"name": "deformation_registration", "domain": "geometry", "description": "3D coordinate deformation regression with a smooth twist and radial compression; exposes spatial coordinates and deformation gradients."}

def _map(x, amp=1.0):
    z=x[:,2]; th=amp*1.8*z
    c=np.cos(th); s=np.sin(th)
    xx=.68*x[:,0]; yy=.68*x[:,1]
    return np.stack([c*xx-s*yy, s*xx+c*yy, z], axis=1).astype(np.float32)

def get_dataset(seed, n_train=400, n_test=200):
    rng=np.random.default_rng(seed)
    xtr=rng.uniform(-1,1,(n_train,3)).astype(np.float32)
    xte=np.random.default_rng(seed+10000).uniform(-1,1,(n_test,3)).astype(np.float32)
    # Fixed moderately large deformation makes the structural regularizer relevant.
    return {"xtr":xtr,"ytr":_map(xtr,2.0),"xte":xte,"yte":_map(xte,2.0),
            "task":"regression","metric":"mse","input_shape":(3,),"out_dim":3}