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

SEED=17
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda': torch.zeros(1,device=device)
except Exception:
    device=torch.device('cpu')

def math_check():
    ps=np.linspace(.01,.99,99); tau=.37; dy=2.3
    rng=np.random.default_rng(SEED); vars=[]
    for p in ps:
        y=np.where(rng.random(200000)<p,dy,0.)
        vars.append(np.var(y))
    vars=np.asarray(vars); theory=ps*(1-ps)*dy**2
    b=np.linspace(-8,8,1001); pp=1/(1+np.exp(-b/tau)); db=b[1]-b[0]
    deriv=np.gradient(pp,db)
    return {'max_abs_variance_error':float(np.max(np.abs(vars-theory))),
            'variance_peak_p':float(ps[np.argmax(theory)]),
            'susceptibility_peak_p':0.5,
            'finite_diff_peak_p':float(pp[np.argmax(deriv)]),
            'susceptibility_bound':1/(4*tau),
            'finite_diff_max':float(deriv.max()),
            'peak_error_from_half':float(abs(ps[np.argmax(theory)]-.5))}

def data(n,T=28,rho=.5,seed=0):
    r=np.random.default_rng(seed); x=np.zeros((n,T+1,1),np.float32); z=np.zeros((n,T),np.float32)
    for i in range(n):
        z[i,0]=r.random()<rho; x[i,0]=r.normal(0,.8)
        for t in range(T):
            if t and r.random()>.94: z[i,t]=r.random()<rho
            a=.82 if z[i,t]==0 else -.72
            x[i,t+1]=a*x[i,t]+r.normal(0,.12)
    return torch.tensor(x),torch.tensor(z)

class Baseline(nn.Module):
    def __init__(self,h=20):
        super().__init__(); self.r=nn.GRU(1,h,batch_first=True); self.out=nn.Linear(h,1)
    def forward(self,x):
        h,_=self.r(x); return self.out(h),h

class ThermoGate(nn.Module):
    def __init__(self,h=20,tau=.65):
        super().__init__(); self.r=nn.GRU(1,h,batch_first=True); self.q=nn.Linear(h,1)
        self.f0=nn.Linear(h,1); self.f1=nn.Linear(h,1); self.tau=tau
    def forward(self,x):
        h,_=self.r(x); q=self.q(h).squeeze(-1); p=torch.sigmoid(q/self.tau)
        y0=self.f0(h).squeeze(-1); y1=self.f1(h).squeeze(-1)
        return (1-p)*y0+p*y1,p,q

def train(model,gated,steps=650):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3); model.train()
    for step in range(steps):
        x,z=data(96,28,.5,1000+step); x=x.to(device); target=x[:,1:,0]
        pred,*rest=model(x[:,:-1]); loss=(pred.squeeze(-1)-target).pow(2).mean()
        if gated:
            p=rest[0]; loss=loss+.015*(p[:,1:]-p[:,:-1]).pow(2).mean()+.003*(p.mean()-.5).pow(2)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),2.0); opt.step()
    return float(loss)

def evaluate(model,gated,rho):
    model.eval(); vals=[]; ps=[]; zs=[]
    with torch.no_grad():
        for k in range(12):
            x,z=data(96,28,rho,50000+int(rho*1000)+k); x=x.to(device); out=model(x[:,:-1])
            vals.append((out[0].squeeze(-1)-x[:,1:,0]).pow(2).mean().item())
            if gated: ps.append(out[1].cpu().numpy()); zs.append(z.numpy())
    ret={'mse':float(np.mean(vals))}
    if gated:
        pp=np.concatenate(ps); zz=np.concatenate(zs)
        ret.update({'mean_gate':float(pp.mean()),'gate_variance':float(pp.var()),
                    'gate_smoothness':float(np.mean(np.diff(pp,axis=1)**2)),
                    'gate_regime_corr':float(np.corrcoef(pp.ravel(),zz.ravel())[0,1])})
    return ret

def run_benchmark():
    global device
    for attempt in range(2):
        try:
            b=Baseline(); g=ThermoGate(); train(b,False); train(g,True)
            result={'params':{'baseline':sum(p.numel() for p in b.parameters()),'idea':sum(p.numel() for p in g.parameters())}}
            for rho in [.1,.5,.9]: result[f'rho_{rho}']={'baseline':evaluate(b,False,rho),'idea':evaluate(g,True,rho)}
            return result
        except RuntimeError as e:
            if attempt==0 and device.type=='cuda':
                device=torch.device('cpu'); torch.cuda.empty_cache(); continue
            raise

def main():
    t=time.time(); results={'device_requested':('cuda' if torch.cuda.is_available() else 'cpu'),'math_check':math_check()}
    results.update(run_benchmark()); results['device_used']=str(device); results['seconds']=time.time()-t
    Path('results.json').write_text(json.dumps(results,indent=2)); print(json.dumps(results,indent=2))
if __name__=='__main__': main()