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

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

# Three mechanism checks, directly tied to the proposed update and halting equations.
def mechanism_checks():
    alpha=0.08; eps=1e-3; n=30
    # For s_(t+1)=(1+alpha*lambda)s_t, boundedness is |1+alpha*lambda|<=1,
    # with strict contraction for <1. The two neutral boundaries are lambda=0,-2/alpha.
    stability=[]
    for lam in [-26,-25,-20,-13,-12.5,-10,0,5]:
        q=1+alpha*lam; vals=np.array([q**t for t in range(n+1)])
        pred='contract' if abs(q)<1 else ('neutral' if abs(q)==1 else 'diverge')
        obs='diverge' if max(abs(vals))>1+1e-8 else ('neutral' if abs(q)==1 else 'contract')
        stability.append({'lambda':lam,'q':q,'predicted':pred,'observed':obs,'abs_s30':float(abs(vals[-1]))})
    # Formula T=min{t: sum_{k=1}^t h_k >= 1-eps}; for constant h this is ceil((1-eps)/h).
    hrows=[]
    for h in [.05,.1,.2,.4,.7]:
        pred=math.ceil((1-eps)/h)
        cumulative=0.; t=0
        while cumulative < 1-eps:
            t+=1; cumulative+=h
        # ACT clipping reaches mass exactly one on the same final update, but its
        # stopping criterion is mass after clipping; report it separately.
        mass=0.; act_t=0
        while mass < 1-eps:
            act_t+=1; mass+=min(h,1-mass)
        hrows.append({'h':h,'formula_T':pred,'observed_formula_T':t,'act_clipped_T':act_t})
    # Since h=sigmoid(b), the predicted T must be non-increasing in b.
    brows=[]
    for b in [-3,-1,0,1,3]:
        h=1/(1+math.exp(-b)); pred=math.ceil((1-eps)/h)
        cumulative=0.; t=0
        while cumulative < 1-eps: t+=1; cumulative+=h
        brows.append({'b':b,'h':h,'formula_T':pred,'observed_T':t})
    return {'stability':stability,'strict_stability_interval':[-25.0,0.0],
            'halting_constant_h':hrows,'halting_bias_sweep':brows}

class ResidualMLP(nn.Module):
    def __init__(self,d=32,blocks=6):
        super().__init__(); self.inp=nn.Linear(2,d); self.blocks=nn.ModuleList([nn.Sequential(nn.LayerNorm(d),nn.Linear(d,d),nn.GELU(),nn.Linear(d,d)) for _ in range(blocks)]); self.out=nn.Linear(d,2)
    def forward(self,x):
        s=self.inp(x)
        for b in self.blocks: s=s+b(s)
        return self.out(s)

class AdaptiveRecurrence(nn.Module):
    def __init__(self,d=32,tmax=8):
        super().__init__(); self.inp=nn.Linear(2,d); self.norm=nn.LayerNorm(d); self.rule=nn.Sequential(nn.Linear(d,48),nn.GELU(),nn.Linear(48,d)); self.halt=nn.Linear(d,1); self.out=nn.Linear(d,2); self.tmax=tmax
    def forward(self,x):
        s=self.inp(x); acc=torch.zeros_like(s); mass=torch.zeros(x.shape[0],1,device=x.device); steps=torch.zeros_like(mass)
        for _ in range(self.tmax):
            s=s+0.25*self.rule(self.norm(s)); h=torch.sigmoid(self.halt(s)); delta=torch.minimum(h,1-mass); acc=acc+delta*s; mass=mass+delta; steps=steps+(mass<1-1e-3).float()
        acc=acc+(1-mass)*s
        return self.out(acc), steps.squeeze(1)

def make_data(n, seed):
    g=np.random.default_rng(seed); x=g.uniform(-1,1,(n,2)).astype('float32'); y=((x[:,0]**2+x[:,1]**2 + .22*x[:,0])>.42).astype('int64')
    return torch.tensor(x),torch.tensor(y)

def train(model, steps=500):
    model.to(device); x,y=make_data(2048,SEED+1); xv,yv=make_data(2048,SEED+2); x,y,xv,yv=[z.to(device) for z in (x,y,xv,yv)]
    opt=torch.optim.Adam(model.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss(); t0=time.perf_counter(); model.train()
    for i in range(steps):
        ix=torch.randint(0,len(x),(128,),device=device); xb,yb=x[ix],y[ix]; opt.zero_grad(); out=model(xb); logits=out[0] if isinstance(out,tuple) else out; loss=lossfn(logits,yb)
        if isinstance(out,tuple): loss=loss+0.002*out[1].mean()
        loss.backward(); opt.step()
    if device=='cuda': torch.cuda.synchronize()
    elapsed=time.perf_counter()-t0; model.eval()
    with torch.no_grad():
        out=model(xv); logits=out[0] if isinstance(out,tuple) else out; acc=(logits.argmax(1)==yv).float().mean().item(); avgsteps=out[1].mean().item() if isinstance(out,tuple) else 6.0
    return {'accuracy':acc,'avg_microsteps':avgsteps,'train_seconds':elapsed,'parameters':sum(p.numel() for p in model.parameters())}

def main():
    result={'device':device,'mechanism_checks':mechanism_checks(),'mini_experiment':{'baseline':train(ResidualMLP()),'idea':train(AdaptiveRecurrence())}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()