import numpy as np
from scipy.linalg import expm


def discretize(A, B, dt):
    n, k = A.shape[0], B.shape[1]
    aug = np.zeros((n + k, n + k)); aug[:n, :n] = A; aug[:n, n:] = B
    E = expm(aug * dt)
    return E[:n, :n], E[:n, n:]


def simulate(m=2, p=2, N=5000, dt=.03, noise=0., seed=0, warmup=1000):
    rng = np.random.default_rng(seed); n = 2
    A = np.array([[-.35, .8], [-.6, -.55]])
    Bbase = np.array([[1.0, .25], [.15, 1.0]])
    B = Bbase[:, :m]
    Cbase = np.array([[1., .25], [-.35, 1.], [.4, -.7]])
    C = Cbase[:p]
    Af = np.array([[0., 1.], [-2.25, -1.5]]); bf = np.array([[0.], [1.]])
    Ad, Bd = discretize(A, B, dt); Df, Gf = discretize(Af, bf, dt)
    x = np.zeros(n); zu = np.zeros((m, n)); zy = np.zeros((p, n)); rows=[]; ys=[]
    for t in range(N + warmup):
        u = .8*rng.normal(size=m) + .35*np.sin(.011*t + np.arange(m))
        y = C @ x + noise*rng.normal(size=p)
        zu = (Df @ zu.T).T + u[:, None]*Gf[:, 0][None, :]
        zy = (Df @ zy.T).T + y[:, None]*Gf[:, 0][None, :]
        x = Ad @ x + Bd @ u
        if t >= warmup: rows.append(np.r_[zu.ravel(), zy.ravel()]); ys.append(y)
    return np.asarray(rows), np.asarray(ys)


def svd_info(Z, tau=1e-3):
    s = np.linalg.svd(Z-Z.mean(0), compute_uv=False)
    return s, int(np.sum(s/s[0] > tau))


def features(H):
    return np.column_stack([H[:,i]*H[:,j] for i in range(H.shape[1]) for j in range(i,H.shape[1])])


def condition(X):
    s=np.linalg.svd(X,compute_uv=False)
    return float(s[0]/max(s[-1],1e-30))


def main():
    np.set_printoptions(precision=4,suppress=True)
    print('PREDICTION 1: SIMO 2n=4; MIMO n(m+1)=6')
    for m,p,label,expected in [(1,2,'SIMO',4),(2,2,'MIMO',6)]:
        Z,_=simulate(m=m,p=p,seed=3); s,r=svd_info(Z)
        print(f'{label}: ambient={Z.shape[1]}, predicted={expected}, observed={r}, ratios={s/s[0]}')

    print('\nPREDICTION 2: filter transient decays and rank approaches theorem')
    for warmup in [0,10,100,300,1000,3000]:
        Z,_=simulate(m=2,p=2,seed=4,warmup=warmup); s,r=svd_info(Z)
        print(f'warmup={warmup:4d}, rank(tau=1e-3)={r}, s7/s1={s[6]/s[0]:.3e}, s8/s1={s[7]/s[0]:.3e}')

    print('\nPREDICTION 2: observation-noise tail scales linearly with noise')
    for noise in [0., 1e-5, 1e-4, 1e-3, 1e-2]:
        Z,_=simulate(m=2,p=2,seed=5,noise=noise); s,_=svd_info(Z)
        tail=float(np.mean(s[6:]/s[0]))
        print(f'noise={noise:.0e}, tail_mean/s1={tail:.3e}, tail/noise={tail/max(noise,1e-30):.3e}')

    print('\nPREDICTION 3: cutoff transition when tau crosses noisy tail')
    Z,_=simulate(m=2,p=2,seed=5,noise=1e-3); s,_=svd_info(Z)
    print('singular_ratios=',s/s[0])
    for tau in [1e-1,1e-2,1e-3,1e-4,1e-5,1e-6,1e-8]:
        _,r=svd_info(Z,tau); print(f'tau={tau:.0e}, rank={r}')

    print('\nCRITIC REGRESSION: quadratic value target, full vs reduced vs under-compressed')
    Z,Y=simulate(m=2,p=2,seed=7,noise=1e-3); split=len(Z)//2; Zc=Z-Z[:split].mean(0)
    U=np.linalg.svd(Zc[:split],full_matrices=False)[2].T
    target=np.sum(Y**2,axis=1)
    for r in [3,6,8]:
        H=Zc@U[:,:r]; X=features(H); tr,te=X[:split],X[split:]
        beta=np.linalg.solve(tr.T@tr+1e-5*np.eye(tr.shape[1]),tr.T@target[:split])
        mse=np.mean((te@beta-target[split:])**2)
        print(f'r={r:2d}, features={X.shape[1]:3d}, test_mse={mse:.5e}, feature_cond={condition(tr):.3e}')

if __name__=='__main__': main()