import json, math, random
from pathlib import Path
import numpy as np
from scipy.integrate import solve_ivp

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


def apply_level(z, j, d, F0, F1, F2, include_quad=True):
    """Derivative of tensor level j, using z[j-1],z[j],z[j+1]."""
    out = np.zeros((d,) * j)
    # z is a list indexed by tensor level; z[0] is scalar
    for oi in np.ndindex(out.shape):
        val = 0.0
        for r in range(j):
            # linear term
            for a in range(d):
                ii = list(oi); ii[r] = a
                val += F1[oi[r], a] * z[j][tuple(ii)]
            # forcing term (for j=1, z[0]=1)
            if j == 1:
                val += F0[oi[r]]
            else:
                ii = tuple(oi[k] for k in range(j) if k != r)
                val += F0[oi[r]] * z[j-1][ii]
            # quadratic term: split the source factor into two indices
            if include_quad:
                for a in range(d):
                    for b in range(d):
                        ii = list(oi); ii[r:r+1] = [a, b]
                        val += F2[oi[r], a*d+b] * z[j+1][tuple(ii)]
        out[oi] = val
    return out


def unpack(y, d, N):
    levels=[np.array(1.0)]
    p=0
    for j in range(1,N+1):
        n=d**j; levels.append(y[p:p+n].reshape((d,)*j)); p += n
    return levels


def pack(levels): return np.concatenate([np.asarray(x).ravel() for x in levels[1:]])
def pack_derivs(levels): return np.concatenate([np.asarray(x).ravel() for x in levels])


def rhs_truncated(t, y, d, N, F0, F1, F2):
    z=unpack(y,d,N); dz=[]
    for j in range(1,N+1):
        # append a harmless zero high level: omitted interaction is precisely excluded
        zz=z + [np.zeros((d,)*(N+1))]
        dz.append(apply_level(zz,j,d,F0,F1,F2,include_quad=(j < N)))
    return pack_derivs(dz)


def rhs_full(t, y, d, Nfull, F0, F1, F2):
    z=unpack(y,d,Nfull); dz=[]
    for j in range(1,Nfull+1):
        zz=z + [np.zeros((d,)*(Nfull+1))]
        dz.append(apply_level(zz,j,d,F0,F1,F2,include_quad=(j < Nfull)))
    return pack_derivs(dz)


def verify_math():
    d=2; N=2; full=3
    F1=np.array([[-0.45,0.12],[-0.08,-0.30]])
    F2=np.array([[0.10,-0.035,0.02,0.015],[-0.04,0.025,-0.02,0.03]])
    F0=np.array([0.08,-0.04]); u0=np.array([0.35,-0.25])
    init=[np.array(1.),u0,np.einsum('i,j->ij',u0,u0),np.einsum('i,j,k->ijk',u0,u0,u0)]
    ts=np.linspace(0,4,401)
    sol=solve_ivp(lambda t,y: rhs_full(t,y,d,full,F0,F1,F2),[0,4],pack(init),t_eval=ts,rtol=1e-10,atol=1e-12)
    tr=solve_ivp(lambda t,y: rhs_truncated(t,y,d,N,F0,F1,F2),[0,4],pack(init[:N+1]),t_eval=ts,rtol=1e-10,atol=1e-12)
    exact=sol.y.T; approx=tr.y.T
    errors=[]; defects=[]; residuals=[]
    dt=ts[1]-ts[0]
    for k in range(len(ts)):
        zfull=unpack(exact[k],d,full); zn=unpack(approx[k],d,N)
        errors.append(np.linalg.norm(exact[k][:d+d*d]-approx[k]))
        defects.append(np.linalg.norm(apply_level(zfull,N,d,F0,F1,F2,include_quad=True)-apply_level(zfull,N,d,F0,F1,F2,include_quad=False)))
        if 1 <= k < len(ts)-1:
            eta=(exact[k+1][:d+d*d]-exact[k-1][:d+d*d])/(2*dt)
            Aeta=rhs_truncated(ts[k], exact[k][:d+d*d],d,N,F0,F1,F2)-rhs_truncated(ts[k],approx[k],d,N,F0,F1,F2)
            residuals.append(np.linalg.norm(eta-Aeta-defects[k]*np.r_[np.zeros(d), np.ones(d*d)]*0))
    # residual above cannot embed the defect into all coordinates; compute exact vector defect
    residuals=[]
    for k in range(len(ts)):
        zfull=unpack(exact[k],d,full); zn=unpack(approx[k],d,N)
        full_proj=rhs_full(ts[k],exact[k],d,full,F0,F1,F2)[:d+d*d]
        trunc_on_exact=rhs_truncated(ts[k],exact[k][:d+d*d],d,N,F0,F1,F2)
        trunc_on_approx=rhs_truncated(ts[k],approx[k],d,N,F0,F1,F2)
        high=apply_level(zfull,N,d,F0,F1,F2,include_quad=True)-apply_level(zfull,N,d,F0,F1,F2,include_quad=False)
        defect=np.r_[np.zeros(d),high.ravel()]
        residuals.append(np.linalg.norm((full_proj-trunc_on_approx)-((trunc_on_exact-trunc_on_approx)+defect)))
    return dict(max_lift_error=float(max(errors)), max_defect=float(max(defects)),
                mean_defect=float(np.mean(defects)), residual_rms=float(np.sqrt(np.mean(np.square(residuals)))),
                final_lift_error=float(errors[-1]))

# Small practical check: learn a forced quadratic sequence, comparing GRU and a lifted Euler cell.
import torch
import torch.nn as nn

def make_data(n, T, d):
    rng=np.random.RandomState(SEED+3)
    xs=rng.randn(n,T,1).astype('float32')
    ys=np.zeros((n,T,d),dtype='float32'); ys[:,0]=rng.randn(n,d)*.2
    F1=np.array([[.82,.10],[-.08,.76]],dtype='float32')[:d,:d]
    F2=np.array([[.16,-.05,.04,.02],[-.07,.05,-.03,.04]],dtype='float32')[:d,:d*d]
    for t in range(T-1):
        u=ys[:,t]; q=np.einsum('bi,bj->bij',u,u).reshape(n,d*d)
        ys[:,t+1]=u@F1.T + q@F2.T + xs[:,t,:,None].squeeze(1)*np.array([.12,-.09],dtype='float32')[:d]
    return torch.tensor(xs),torch.tensor(ys)

class GRUModel(nn.Module):
    def __init__(self,d=2,h=8):
        super().__init__(); self.r=nn.GRU(1,h,batch_first=True); self.o=nn.Linear(h,d)
    def forward(self,x): return self.o(self.r(x)[0])

class LiftModel(nn.Module):
    """N=2 Carleman cell; z2 is maintained explicitly and contractions are factor-wise."""
    def __init__(self,d=2,dt=0.02):
        super().__init__(); self.d=d; self.dt=dt
        self.f0=nn.Linear(1,d); self.f1=nn.Linear(1,d*d); self.f2=nn.Linear(1,d*d*d)
        for layer in (self.f0,self.f1,self.f2):
            nn.init.normal_(layer.weight, 0.0, 0.08); nn.init.zeros_(layer.bias)
    def forward(self,x):
        B,T,_=x.shape; d=self.d; h=torch.zeros(B,d,device=x.device)
        z2=h[:,:,None]*h[:,None,:]; outs=[]; defects=[]
        for t in range(T):
            f0=self.f0(x[:,t]); f1=self.f1(x[:,t]).view(B,d,d); f2=self.f2(x[:,t]).view(B,d,d,d)
            z3=h[:,:,None,None]*z2[:,None,:,:]
            # omitted A_3^2 z3, summed over the two tensor-factor positions
            quad2=torch.einsum('b i a c,b a c j->b i j',f2,z3)
            quad2=quad2+torch.einsum('b j a c,b i a c->b i j',f2,z3)
            # equivalent first-level quadratic contribution F2 z2
            dh=torch.einsum('b i a,b a->b i',f1,h)+f0+torch.einsum('b i a c,b a c->b i',f2,z2)
            dz2=quad2+torch.einsum('b i a,b a j->b i j',f1,z2)+torch.einsum('b j a,b i a->b i j',f1,z2)
            dz2=dz2+f0[:,:,None]*h[:,None,:]+h[:,:,None]*f0[:,None,:]
            defects.append(torch.linalg.vector_norm(quad2.reshape(B,-1),dim=1))
            h=h+self.dt*dh; z2=z2+self.dt*dz2; outs.append(h)
        return torch.stack(outs,1),torch.stack(defects,1)

def train(model,x,y,steps=350):
    requested = "cuda" if torch.cuda.is_available() else "cpu"
    def run(device):
        model.to(device); xx=x.to(device); yy=y.to(device)
        opt=torch.optim.Adam(model.parameters(),lr=1e-3)
        model.train()
        for _ in range(steps):
            opt.zero_grad()
            out=model(xx)[0] if isinstance(model,LiftModel) else model(xx)
            loss=((out-yy)**2).mean()
            if not torch.isfinite(loss): raise FloatingPointError('non-finite training loss')
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
        model.eval()
        with torch.no_grad():
            out=model(xx)[0] if isinstance(model,LiftModel) else model(xx)
        return float(((out-yy)**2).mean().cpu())
    try:
        return run(requested)
    except Exception as exc:
        if requested != "cuda": raise
        print("CUDA failed; retrying on CPU:", repr(exc))
        return run("cpu")

def main():
    check=verify_math(); torch.manual_seed(SEED)
    x,y=make_data(96,32,2)
    b=train(GRUModel(),x,y); torch.manual_seed(SEED); i=train(LiftModel(),x,y)
    result={'math_check':check,'sequence_train_mse':{'baseline_gru':b,'carleman_lift':i}}
    Path('results.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))
if __name__=='__main__': main()