import json, numpy as np
from vdp_cell import rk4_step

rng=np.random.default_rng(11)
N=700; d=16; delay=80; wash=20
# Binary delayed-memory task: predict input from the state at the current time.
u=rng.uniform(-1,1,size=N)
target=u.copy()
B=rng.normal(0,.7,size=(d,))
W=rng.normal(0,1/np.sqrt(d),size=(d,d)); W/=max(1.0,np.linalg.eigvals(W).real.max()/0.9)

def ridge(X,y,l=1e-3):
    A=X.T@X+l*np.eye(X.shape[1]); return np.linalg.solve(A,X.T@y)
def score(X,y):
    cut=int(.65*len(y)); w=ridge(X[:cut],y[:cut]); pred=X[cut:]@w
    return float(np.mean((pred-y[cut:])**2)),float(np.corrcoef(pred,y[cut:])[0,1])
# Same scalar drive mapped into every cell; compare tanh Elman and VdP oscillator.
for kind in ['tanh','vdp']:
    h=np.zeros(d); feats=[]
    for t in range(N):
        if kind=='tanh': h=np.tanh(W@h+B*u[t])
        else:
            drive=np.zeros(d); drive[:d//2]=B[:d//2]*u[t]; drive[d//2:]=B[d//2:]*u[t]
            h=rk4_step(h,.05,np.full(d//2,2.),1.,1.,drive)
        # state must carry a delay; train target at t-delay from current state
        if t>=delay: feats.append((h.copy(),target[t-delay]))
    X=np.array([x for x,y in feats]); y=np.array([y for x,y in feats])
    mse,corr=score(np.c_[X,np.ones(len(X))],y)
    print(kind, {'test_mse':mse,'test_corr':corr})