import json, os, random
import numpy as np

SEED = 2139
np.random.seed(SEED); random.seed(SEED)


def parts(V):
    return (V + V.T) / 2.0, (V - V.T) / 2.0


def lift(V, m=4):
    d = V.shape[0]
    S, A = parts(V)
    L = np.zeros((m*d, m*d))
    for i in range(m):
        L[i*d:(i+1)*d, i*d:(i+1)*d] = A
        if i + 1 < m:
            L[i*d:(i+1)*d, (i+1)*d:(i+2)*d] = S
            L[(i+1)*d:(i+2)*d, i*d:(i+1)*d] = -S
    return L, S, A


def embed(x, m):
    z = np.zeros(m * len(x)); z[:len(x)] = x; z[len(x):2*len(x)] = x
    return z


def cayley(L, eta):
    I = np.eye(len(L))
    return np.linalg.solve(I - .5*eta*L, I + .5*eta*L)


def run_checks():
    # Fixed non-normal, damped, transient-amplifying generator.
    V0 = np.array([[-.35, 2.0], [-.15, -.10]])
    x = np.array([.8, -.4])
    m = 5
    L, S, A = lift(V0, m)
    skew_err = np.linalg.norm(L + L.T, ord=np.inf)
    proj_err = np.linalg.norm((L[:2,:] @ embed(x,m)) - V0 @ x)
    z0 = embed(x,m)
    cayley_rows = []
    euler_rows = []
    for gamma in [0.0, .25, .5, 1.0, 2.0]:
        V = gamma * V0
        Lg, _, _ = lift(V,m)
        lam = np.linalg.norm(Lg, 2)
        for eta in [.05, .2, .5]:
            z = embed(x,m)
            C = cayley(Lg,eta)
            # ten steps: predicted exact preservation for every eta and gamma
            for _ in range(10): z = C @ z
            rel_norm_error = abs(np.linalg.norm(z)/np.linalg.norm(embed(x,m))-1)
            cayley_rows.append([gamma, eta, rel_norm_error])
        # Euler one-step identity and accumulated growth at eta=.2
        eta=.2; z=embed(x,m); z1=(np.eye(len(Lg))+eta*Lg)@z
        exact_increment = np.linalg.norm(z1)**2-np.linalg.norm(z)**2
        predicted_increment = eta**2*np.linalg.norm(Lg@z)**2
        for _ in range(10): z=(np.eye(len(Lg))+eta*Lg)@z
        euler_rows.append([gamma, exact_increment, predicted_increment,
                           np.linalg.norm(z)/np.linalg.norm(embed(x,m))])

    # Projected transient gain sweep: ||P exp(tL) E|| / ||E|| <= 1.
    gain_rows=[]
    from scipy.linalg import expm
    for gamma in [0.0, .25, .5, 1., 2.]:
        Lg,_,_=lift(gamma*V0,m)
        gains=[]
        for t in np.linspace(0,4,81):
            gains.append(np.linalg.norm(expm(t*Lg)[:2,:2],2) / np.sqrt(2))
        gain_rows.append([gamma, float(max(gains)), float(np.sqrt(2)), float(gains[-1])])
    return {
        'skew_inf_error': float(skew_err),
        'projection_derivative_error': float(proj_err),
        'cayley_rows_gamma_eta_rel_norm_error': cayley_rows,
        'euler_rows_gamma_increment_predicted_increment_norm_ratio_after_10': euler_rows,
        'projected_gain_rows_gamma_max_gain_bound_final_gain': gain_rows,
    }


def tiny_task():
    # A matched-parameter linear memory task, trained with torch if available.
    try:
        import torch
        torch.manual_seed(SEED)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if device == 'cuda':
            try: torch.cuda.set_per_process_memory_fraction(.15)
            except Exception: pass
        T, N, d = 40, 256, 4
        rng=np.random.default_rng(SEED)
        u=torch.tensor(rng.normal(size=(N,T,2)),dtype=torch.float32,device=device)
        # target is a stable but non-normal recurrence, final-state regression
        V=torch.tensor([[-.18,1.2,0,0],[-.05,-.12,0,0],[0,0,-.3,.7],[0,0,-.1,-.25]],dtype=torch.float32,device=device)
        B=torch.tensor(rng.normal(scale=.4,size=(d,2)),dtype=torch.float32,device=device)
        target=torch.zeros(N,d,device=device)
        for k in range(T): target=target + .08*(target@V.T + u[:,k]@B.T)
        target=target.detach()
        # baseline vanilla RNN, small and deliberately matched by hidden dimension
        class Base(torch.nn.Module):
            def __init__(self):
                super().__init__(); self.W=torch.nn.Parameter(torch.randn(d,d,device=device)*.15); self.U=torch.nn.Parameter(torch.randn(d,2,device=device)*.15); self.q=torch.nn.Parameter(torch.randn(d,d,device=device)*.15)
            def forward(self, x):
                h=torch.zeros(x.shape[0],d,device=device)
                for k in range(T): h=torch.tanh(h@self.W.T+x[:,k]@self.U.T)
                return h@self.q.T
        # lift has 2d state and linear readout; stable Cayley, train only input/readout
        class Lift(torch.nn.Module):
            def __init__(self):
                super().__init__(); self.rawS=torch.nn.Parameter(torch.randn(d,d,device=device)*.12); self.rawA=torch.nn.Parameter(torch.randn(d,d,device=device)*.12); self.U=torch.nn.Parameter(torch.randn(2*d,2,device=device)*.15); self.q=torch.nn.Parameter(torch.randn(d,d,device=device)*.15)
            def forward(self,x):
                S=self.rawS+self.rawS.T; A=self.rawA-self.rawA.T; z=torch.zeros(x.shape[0],2*d,device=device); eta=.08
                L=torch.cat([torch.cat([A,S],1),torch.cat([-S,A],1)],0); I=torch.eye(2*d,device=device); C=torch.linalg.solve(I-eta*L/2,I+eta*L/2)
                for k in range(T): z=z@C.T + x[:,k]@self.U.T
                return z[:,:d]@self.q.T
        def fit(model):
            opt=torch.optim.Adam(model.parameters(),lr=.02); losses=[]
            for _ in range(120):
                pred=model(u); loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); opt.step(); losses.append(float(loss.detach().cpu()))
            return losses[0], losses[-1]
        b=fit(Base()); l=fit(Lift())
        return {'device':device,'baseline_initial_final_mse':b,'lift_initial_final_mse':l}
    except Exception as e:
        return {'error':repr(e)}

if __name__ == '__main__':
    out={'checks':run_checks(),'tiny_task':tiny_task()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))