import numpy as np
META = {"name":"correlated_multitask_regression","domain":"multi_task_learning","description":"Six regression tasks with latent task groups and correlated observation residuals; used to test residual-redundancy routing."}

def get_dataset(seed, n_train, n_test):
    rng=np.random.default_rng(seed); p=12; t=6
    w=rng.normal(size=(3,p)).astype(np.float32)
    nuisance=rng.normal(size=(t,p)).astype(np.float32)
    cov=np.eye(t,dtype=np.float32)*.20
    for a,b in [(0,1),(2,3),(4,5)]: cov[a,b]=cov[b,a]=.16
    def make(n):
        x=rng.normal(size=(n,p)).astype(np.float32)
        y=np.zeros((n,t),dtype=np.float32)
        for j,g in enumerate([[0,1],[2,3],[4,5]]):
            latent=np.tanh(x@w[j]) + .35*np.sin(x[:,j%4])
            for q in g: y[:,q]=latent + .12*(x@nuisance[q])
        y += rng.multivariate_normal(np.zeros(t),cov,size=n).astype(np.float32)
        return x,y
    xtr,ytr=make(n_train); xte,yte=make(n_test)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":6}