import json
import numpy as np
from scipy.optimize import brentq


def transition(g, taps):
    taps=np.asarray(taps,float); T=np.zeros((1+len(taps),1+len(taps)))
    T[0,0]=1.; T[0,1:]=g*taps; T[1:,:-1]=np.eye(len(taps)); return T


def margin(g,taps,sigma=1,window=None):
    if not np.isfinite(g): return float('nan')
    taps=np.asarray(taps,float); window=len(taps)+1 if window is None else window
    R=np.linalg.matrix_power(transition(g,taps),window)-sigma*np.eye(len(taps)+1)
    return float(np.linalg.svd(R,compute_uv=False)[-1])


def determinant(g,taps,sigma):
    T=transition(g,taps); n=T.shape[0]
    return float(np.linalg.det(np.linalg.matrix_power(T,n)-sigma*np.eye(n)))


def roots_in_range(taps,sigma,lo=-8.,hi=8.):
    xs=np.linspace(lo,hi,32001); ys=np.array([determinant(x,taps,sigma) for x in xs]); roots=[]
    for i in range(len(xs)-1):
        if not (np.isfinite(ys[i]) and np.isfinite(ys[i+1])): continue
        if ys[i]==0: roots.append(xs[i])
        elif ys[i]*ys[i+1]<0:
            roots.append(brentq(lambda z:determinant(z,taps,sigma),xs[i],xs[i+1]))
    return sorted(set(round(float(r),12) for r in roots))


def first_nontrivial_root(taps,sigma):
    roots=[r for r in roots_in_range(taps,sigma) if abs(r)>1e-4]
    return (min(roots,key=abs),'determinant_root') if roots else (float('nan'),'none')


def grid_minimum(taps,sigma,lo=-8.,hi=8.):
    xs=np.linspace(lo,hi,32001); ys=np.array([margin(x,taps,sigma) for x in xs])
    i=int(np.nanargmin(ys)); return float(xs[i]),float(ys[i])


def numerical_check():
    base=np.array([.7,-.25,.12]); rp,_=first_nontrivial_root(base,1); rm,km=first_nontrivial_root(base,-1)
    rows=[]
    for a in [.5,1.,2.]:
        r,method=first_nontrivial_root(a*base,1)
        rows.append({'tap_scale':a,'predicted_signed_gain':rp/a,'observed_signed_gain':r,'relative_error':None if not np.isfinite(r) else abs(r-rp/a)/(abs(rp/a)+1e-12),'method':method})
    am,amin=grid_minimum(base,-1)
    return {'predictions':['Fixed tap ratios imply roots scale as 1/a because Phi(g,a*taps)=Phi(a*g,taps).','A true periodic/antiperiodic bifurcation has det(Phi-sigma I)=0 and margin s_min=0.','The sigma sign selects distinct modes; absence of a root is itself a negative result.'],'base_periodic_gain':rp,'base_antiperiodic_gain':None if not np.isfinite(rm) else rm,'periodic_margin_at_root':margin(rp,base,1),'antiperiodic_root_margin':None if not np.isfinite(rm) else margin(rm,base,-1),'antiperiodic_grid_min_gain':am,'antiperiodic_grid_min_margin':amin,'scaling_rows':rows,'periodic_roots':roots_in_range(base,1),'antiperiodic_roots':roots_in_range(base,-1)}


def train_comparison(seed=7):
    import torch
    torch.manual_seed(seed); np.random.seed(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        n,L,delay=96,24,3; X=torch.randn(n,L,1,device=device); Y=torch.zeros_like(X); Y[:,delay:]=X[:,:-delay]
        class Delayed(torch.nn.Module):
            def __init__(self):
                super().__init__(); self.inp=torch.nn.Linear(1,8); self.out=torch.nn.Linear(8,1); self.g=torch.nn.Parameter(torch.tensor(.65)); self.taps=torch.nn.Parameter(torch.tensor([.55,-.25,.10]))
            def forward(self,x):
                h=[torch.zeros(x.shape[0],8,device=x.device) for _ in range(4)]; ys=[]
                for t in range(L):
                    cur=torch.tanh(self.inp(x[:,t])+self.g*(self.taps[0]*h[-1]+self.taps[1]*h[-2]+self.taps[2]*h[-3])); h.append(cur); h=h[-4:]; ys.append(self.out(cur))
                return torch.stack(ys,1)
            def penalty(self):
                T=torch.zeros((4,4),device=device); T[0,0]=1.; T[0,1:]=self.g*self.taps; T[1:,:-1]=torch.eye(3,device=device); P=torch.matrix_power(T,4)
                return sum(torch.relu(.08-torch.linalg.svdvals(P-s*torch.eye(4,device=device))[-1])**2 for s in (1.,-1.))
        results={}
        for name,monitored in [('baseline',False),('idea',True)]:
            model=Delayed().to(device); opt=torch.optim.Adam(model.parameters(),lr=.025); obj=0.
            for _ in range(180):
                task=torch.mean((model(X)-Y)**2); loss=task+(.8*model.penalty() if monitored else 0.); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),2.); opt.step(); obj=float(loss.detach().cpu())
            with torch.no_grad(): task=float(torch.mean((model(X)-Y)**2).cpu()); g=float(model.g.cpu()); t=model.taps.detach().cpu().numpy()
            results[name]={'task_mse':task,'final_objective':obj,'gain':g,'taps':t.tolist(),'periodic_margin':margin(g,t,1),'antiperiodic_margin':margin(g,t,-1)}
        return {'device':device,'results':results}
    except Exception as e: return {'device':'cpu-fallback','error':repr(e)}

if __name__=='__main__':
    out={'math_check':numerical_check(),'training':train_comparison()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))