import json, math, random
import numpy as np

# ISS-Constrained Modular Recurrent Network: toy verification + tiny task comparison.
# The analytic toy recurrence is x_{k+1}=q*x_k+c, q=1+dt*(lam-gamma).

def set_seed(seed=7):
    random.seed(seed); np.random.seed(seed)


def toy_verification():
    dt, lam = 0.1, 2.0
    # Prediction 1: stability boundary is gamma=lam (for positive Euler q),
    # with divergence for gamma < lam and contraction for gamma > lam.
    gammas = np.array([0.0, 1.0, 1.9, 2.0, 2.1, 3.0, 10.0])
    boundary_rows=[]
    for g in gammas:
        q=1+dt*(lam-g)
        x=1.0
        for _ in range(80): x=q*x
        boundary_rows.append({'gamma':float(g),'q_pred':float(q),'abs_q':float(abs(q)),
                              'final_abs':float(abs(x)),'stable_pred':bool(abs(q)<1)})
    # Prediction 2: log perturbation slope equals log |q| in stable regime.
    slope_rows=[]
    for g in [2.1, 3.0, 5.0]:
        q=1+dt*(lam-g); vals=[]; d=1.0
        for _ in range(45):
            vals.append(abs(d)); d=q*d
        slope=np.polyfit(np.arange(8,45), np.log(np.maximum(vals[8:],1e-300)), 1)[0]
        slope_rows.append({'gamma':g,'pred_log_abs_q':float(np.log(abs(q))),
                           'observed_slope':float(slope),'relative_error':float(abs(slope-math.log(abs(q)))/abs(math.log(abs(q))))})
    # Prediction 3: constant forcing reaches c/(gamma-lambda), and geometric bound
    # using measured contraction a=|q| is c*dt/(1-a), exactly equal here.
    force_rows=[]
    c=0.7
    for g in [2.1, 3.0, 5.0]:
        q=1+dt*(lam-g); x=0.0
        for _ in range(500): x=q*x+dt*c
        measured=abs(x); exact=c/(g-lam); bound=dt*c/(1-abs(q))
        force_rows.append({'gamma':g,'observed_limit':float(measured),'predicted_limit':float(exact),
                           'geometric_bound':float(bound),'relative_error':float(abs(measured-exact)/exact)})
    # Perceptual claim: z_{k+1}=alpha*z_k has decay slope log(alpha).
    alpha=0.93; z=1.; zs=[]
    for _ in range(80): zs.append(abs(z)); z*=alpha
    z_slope=np.polyfit(np.arange(10,80),np.log(np.maximum(zs[10:],1e-300)),1)[0]
    return {'boundary_sweep':boundary_rows,'decay_sweep':slope_rows,
            'iss_forcing_sweep':force_rows,
            'perception':{'alpha':alpha,'pred_log_alpha':math.log(alpha),'observed_slope':float(z_slope)}}


def task_comparison():
    # Small sequence regression: predict normalized sum of inputs from the final state.
    try:
        import torch
        import torch.nn as nn
        torch.set_num_threads(4)
        torch.manual_seed(7); np.random.seed(7)
        device=torch.device('cpu')  # deterministic and avoids shared cuDNN allocation failures
        try:
            # Probe CUDA and fall back on any allocation/runtime issue.
            if device.type=='cuda': torch.zeros(1,device=device)
        except Exception: device=torch.device('cpu')
        T,N,B,H,Z=40,1,64,24,12
        class Modular(nn.Module):
            def __init__(self):
                super().__init__(); self.h=H; self.z=Z; self.alpha=.92; self.dt=.1; self.gamma=1.5
                self.pz=nn.Linear(Z,Z,bias=False); self.pu=nn.Linear(N,Z)
                self.fx=nn.Linear(H,H,bias=False); self.fz=nn.Linear(Z,H); self.fu=nn.Linear(N,H); self.head=nn.Linear(H,1)
                nn.init.orthogonal_(self.pz.weight); self.pz.weight.data.mul_(self.alpha)
            def forward(self,u):
                b=u.shape[1]; z=torch.zeros(b,Z,device=u.device); x=torch.zeros(b,H,device=u.device)
                for k in range(u.shape[0]):
                    z=torch.tanh(self.pz(z)+self.pu(u[k]))
                    f=torch.tanh(self.fx(x)+self.fz(z)+self.fu(u[k]))
                    x=x+self.dt*(f-self.gamma*x)
                return self.head(x).squeeze(-1)
        class Vanilla(nn.Module):
            def __init__(self):
                super().__init__(); self.r=nn.GRU(N,H); self.head=nn.Linear(H,1)
            def forward(self,u): return self.head(self.r(u)[0][-1]).squeeze(-1)
        def run(model):
            model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3); losses=[]
            for step in range(90):
                u=torch.randn(T,B,N,device=device); y=u.sum(0).squeeze(-1)/math.sqrt(T)
                opt.zero_grad(); pred=model(u); loss=((pred-y)**2).mean(); loss.backward(); opt.step(); losses.append(float(loss.detach().cpu()))
            with torch.no_grad():
                u=torch.randn(T,256,N,device=device); y=u.sum(0).squeeze(-1)/math.sqrt(T); test=float(((model(u)-y)**2).mean().cpu())
            return {'final_train_mse':losses[-1],'test_mse':test,'parameters':sum(p.numel() for p in model.parameters())}
        try:
            return {'device':str(device),'vanilla_gru':run(Vanilla()),'iss_modular':run(Modular())}
        except Exception as first_error:
            # Shared CUDA environments can fail during cuDNN workspace allocation;
            # rerun from fresh CPU modules, preserving the experiment definition.
            if device.type == 'cuda':
                device=torch.device('cpu')
                return {'device':str(device),'cuda_error':repr(first_error),
                        'vanilla_gru':run(Vanilla()),'iss_modular':run(Modular())}
            raise
    except Exception as e:
        return {'error':repr(e)}

if __name__=='__main__':
    set_seed(7)
    out={'toy':toy_verification(),'task':task_comparison()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))