"""Custom structural track: parameterized clamped-free beam mode frequencies."""
import numpy as np
from scipy.linalg import eigh
META = {"name":"beam_eigenfrequency","domain":"structural_eigenproblem","description":"Predict one clamped-free beam eigenfrequency from stiffness/mass scaling; FEM dynamic stiffness supplies WW brackets."}

def matrices(scale=1.0, nel=4):
    h=1.0/nel; nd=2*(nel+1); K=np.zeros((nd,nd)); M=np.zeros((nd,nd))
    ke=scale/h**3*np.array([[12,6*h,-12,6*h],[6*h,4*h*h,-6*h,2*h*h],[-12,-6*h,12,-6*h],[6*h,2*h*h,-6*h,4*h*h]],float)
    me=h/420*np.array([[156,22*h,54,-13*h],[22*h,4*h*h,13*h,-3*h*h],[54,13*h,156,-22*h],[-13*h,-3*h*h,-22*h,4*h*h]],float)
    for e in range(nel):
        ix=[2*e,2*e+1,2*e+2,2*e+3]; K[np.ix_(ix,ix)]+=ke; M[np.ix_(ix,ix)]+=me
    return K[2:,2:],M[2:,2:]

def sample(seed, n):
    rng=np.random.default_rng(seed); scales=rng.uniform(.75,1.25,n)
    # input is scalar scale; targets are the third mode frequency
    ys=[]
    for s in scales:
        K,M=matrices(float(s)); ev=eigh(K,M,eigvals_only=True); ys.append(np.sqrt(ev[2]))
    return scales.astype(np.float32)[:,None],np.asarray(ys,np.float32)[:,None]

def get_dataset(seed,n_train,n_test):
    xtr,ytr=sample(seed,n_train); xte,yte=sample(seed+100003,n_test)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse"}