import numpy as np

META = {"name": "strain_congruence", "domain": "dynamics", "description": "Volume-preserving 3D deformation with six directional strain channels; predict a nonlinear strain-dynamics target."}


def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng = np.random.RandomState(s)
        X = np.empty((n, 15), np.float32)
        Y = np.empty((n, 1), np.float32)
        phi = (1 + np.sqrt(5.0)) / 2
        axes = np.array([(0,1,phi),(1,phi,0),(phi,0,1),(0,1,-phi),(1,-phi,0),(phi,0,-1)], np.float32) / np.sqrt(1+phi*phi)
        for i in range(n):
            H = rng.normal(0, .65, (3,3)).astype(np.float64)
            H -= np.trace(H) * np.eye(3) / 3
            # diagonal positive SL(3) deformation, preserving the stated structure
            F = np.diag(np.exp(np.diag(H)))
            S = rng.normal(0, .7, (3,3)); S = (S+S.T)/2
            S -= np.trace(S)*np.eye(3)/3
            w = axes @ F.T
            y = np.einsum('ni,ij,nj->n', w, S, w)
            # target is a task-independent nonlinear dynamics quantity
            target = (np.tanh(S[0,1] + .35*S[1,2]) + .25*np.sin(np.linalg.norm(F-np.eye(3))) + .1*S[0,0]*S[1,1])
            X[i] = np.r_[F.reshape(-1), y]
            Y[i,0] = target
        return X, Y
    xtr,ytr=make(n_train,seed); xte,yte=make(n_test,seed+5000)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":1}