import numpy as np
META = {"name":"relativistic_conservative_regression","domain":"conservative_state_admissibility","description":"Predict valid relativistic conservative states from valid baselines and controls; GQL residual limiting is structurally native."}

def _make(seed, n):
    r=np.random.RandomState(seed)
    # Inputs are a valid baseline U0=(D,mx,my,E0) plus a requested physical perturbation.
    d=r.uniform(.4,2.0,n); mx=r.normal(0,.35,n); my=r.normal(0,.35,n)
    e=np.sqrt(d*d+mx*mx+my*my)+r.uniform(.15,.8,n)
    u0=np.stack([d,mx,my,e],1)
    # A smooth target state; training noise makes unconstrained extrapolations common.
    dd=.45*np.tanh(r.normal(size=n)); dm=r.normal(0,.38,(n,2));
    target=u0.copy(); target[:,0]=d+dd
    target[:,1:3]+=dm
    target[:,3]=np.sqrt(target[:,0]**2+np.sum(target[:,1:3]**2,1))+r.uniform(.08,.65,n)
    # The model receives baseline and a control-like requested delta.
    x=np.concatenate([u0, target-u0],1).astype(np.float32)
    return x,target.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":4}