import json, math, time, random
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 3130
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda':
        torch.cuda.empty_cache()
except Exception:
    device = torch.device('cpu')

# ---------- Stage 1: direct numerical verification of the stated recurrence ----------
def verify_math():
    rng = np.random.default_rng(SEED)
    dg, dl, H = 3, 6, 20
    # Stable block operators, as recommended by the idea.
    Ag = np.array([[.97,-.12,.0],[.12,.97,.0],[0,0,.93]])
    Al = np.zeros((dl, dl))
    for i, (a,b) in enumerate([(0.985, .08), (0.96,.18), (0.92,.25)]):
        Al[2*i:2*i+2,2*i:2*i+2] = [[a,-b],[b,a]]
    B = rng.normal(0,.04,(dg+dl,2))
    z0 = rng.normal(size=dg+dl); us = rng.normal(size=(H,2))
    A = np.zeros((dg+dl,dg+dl)); A[:dg,:dg]=Ag; A[dg:,dg:]=Al
    zs = [z0]
    for u in us: zs.append(A @ zs[-1] + B @ u)
    zs = np.asarray(zs)
    # Same update computed as separate channels; this tests the block formula itself.
    zg, zl = z0[:dg].copy(), z0[dg:].copy(); separate=[np.r_[zg,zl]]
    for u in us:
        zg = Ag @ zg + B[:dg] @ u
        zl = Al @ zl + B[dg:] @ u
        separate.append(np.r_[zg,zl])
    recurrence_error = float(np.max(np.abs(zs-np.asarray(separate))))
    # Stability signal: with no controls, the norm contracts because both blocks do.
    no_control = [z0.copy()]
    for _ in range(H): no_control.append(A @ no_control[-1])
    norms = np.linalg.norm(no_control, axis=1)
    contraction_ratio = float(norms[-1]/norms[0])
    spectral_global = float(max(abs(np.linalg.eigvals(Ag))))
    spectral_local = float(max(abs(np.linalg.eigvals(Al))))
    return dict(recurrence_max_abs_error=recurrence_error,
                spectral_radius_global=spectral_global,
                spectral_radius_local=spectral_local,
                zero_input_norm_ratio=contraction_ratio,
                stability_observed=(spectral_global < 1 and spectral_local < 1 and contraction_ratio < 1))

# ---------- Tiny nonlinear observation system with known global/local latent structure ----------
def make_data(n=1100, T=24):
    rng = np.random.default_rng(SEED+1)
    dg, dl = 3, 6
    Ag = np.array([[.975,-.11,.0],[.11,.975,.0],[0,0,.94]], dtype=np.float32)
    Al = np.zeros((dl,dl), dtype=np.float32)
    for i,(a,b) in enumerate([(0.985,.075),(.965,.16),(.93,.22)]):
        Al[2*i:2*i+2,2*i:2*i+2] = [[a,-b],[b,a]]
    # Mild global-to-local physical coupling makes the observation nonlinear while
    # preserving a useful approximately block-structured latent representation.
    xs = np.zeros((n,T,12), np.float32)
    for k in range(n):
        g = rng.normal(0,.8,dg).astype(np.float32)
        l = rng.normal(0,.8,dl).astype(np.float32)
        for t in range(T):
            # nonlinear observation: global coordinates, local coordinates, and products
            prod = (g[0] * l[::2] + g[1] * l[1::2]).astype(np.float32)
            xs[k,t] = np.r_[g, l, prod]
            g = Ag @ g + rng.normal(0,.012,dg).astype(np.float32)
            l = Al @ l + np.repeat(np.tanh(g[:3]),2)[:dl].astype(np.float32)*.018
    return torch.tensor(xs)

class MLP(nn.Module):
    def __init__(self, a,b,c): super().__init__(); self.net=nn.Sequential(nn.Linear(a,b),nn.Tanh(),nn.Linear(b,c))
    def forward(self,x): return self.net(x)

class BlockKoopman(nn.Module):
    def __init__(self):
        super().__init__(); self.eg=MLP(12,32,3); self.el=MLP(12,32,6)
        self.Ag=nn.Parameter(torch.eye(3)*.96); self.Al=nn.Parameter(torch.eye(6)*.96)
        self.dec=MLP(9,40,12)
    def forward(self,x0,H):
        g,l=self.eg(x0),self.el(x0); out=[]
        for _ in range(H):
            g=torch.einsum('bi,ji->bj',g,self.Ag); l=torch.einsum('bi,ji->bj',l,self.Al)
            out.append(self.dec(torch.cat([g,l],1)))
        return torch.stack(out,1)

class GRUBaseline(nn.Module):
    def __init__(self):
        super().__init__(); self.enc=MLP(12,32,12); self.gru=nn.GRUCell(12,12); self.dec=MLP(12,40,12)
    def forward(self,x0,H):
        h=self.enc(x0); out=[]
        for _ in range(H):
            h=self.gru(torch.zeros_like(h),h); out.append(self.dec(h))
        return torch.stack(out,1)

def train_eval(model, train, test, epochs=75):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    xtr=train[:,0].to(device); ytr=train[:,1:11].to(device)
    bs=64; t0=time.perf_counter(); model.train()
    for ep in range(epochs):
        perm=torch.randperm(len(xtr),device=device)
        for ii in range(0,len(xtr),bs):
            ix=perm[ii:ii+bs]; pred=model(xtr[ix],10)
            loss=((pred-ytr[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    if device.type=='cuda': torch.cuda.synchronize()
    train_seconds=time.perf_counter()-t0
    model.eval();
    with torch.no_grad():
        pred=model(test[:,0].to(device),20).cpu(); target=test[:,1:21]
    mse=((pred-target)**2).mean(2).mean(0).numpy()
    return dict(mse_1=float(mse[0]),mse_10=float(mse[9]),mse_20=float(mse[19]),train_seconds=train_seconds,
                params=sum(p.numel() for p in model.parameters())), model

def main():
    math_check=verify_math()
    data=make_data(); train=data[:900]; test=data[900:]
    # Identical initialization order and data; each model gets the same objective/horizon.
    torch.manual_seed(SEED+10); base, _=train_eval(GRUBaseline(),train,test)
    torch.manual_seed(SEED+10); idea, model=train_eval(BlockKoopman(),train,test)
    with torch.no_grad():
        rho_g=max(abs(np.linalg.eigvals(model.Ag.detach().cpu().numpy())))
        rho_l=max(abs(np.linalg.eigvals(model.Al.detach().cpu().numpy())))
    idea['learned_rho_global']=float(rho_g); idea['learned_rho_local']=float(rho_l)
    result={'device':str(device),'math_check':math_check,'baseline':base,'idea':idea,
            'lower_mse_at_20':idea['mse_20'] < base['mse_20']}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()