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

SEED = 1628
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def jacobian(kE, kI):
    return np.array([[2*kE-2, 0, 2*kI],
                     [0, 2*kI-2, 2*kE],
                     [kE, kI, kE+kI-2]], dtype=float)

def predicted_eigs(kE, kI):
    d = kE-kI
    return np.array([-2., 2*(d-1), d-2])

def fixed_point(kE, kI, D):
    return -np.linalg.solve(jacobian(kE,kI), np.asarray(D, dtype=float))

def toy_checks():
    # The supplied J is checked directly. Its characteristic polynomial is
    # (lambda+2)(lambda-2(kE+kI-1))(lambda-(kE+kI-2)).
    max_stated_err = 0.; max_corrected_err = 0.; boundary = []
    for kI in np.linspace(.05,.8,8):
        for d in np.linspace(-.8,1.25,18):
            kE=d+kI; vals=np.sort(np.linalg.eigvals(jacobian(kE,kI)).real)
            stated=np.sort(predicted_eigs(kE,kI))
            s=kE+kI
            corrected=np.sort(np.array([-2.,2*(s-1),s-2]))
            max_stated_err=max(max_stated_err,float(np.max(np.abs(vals-stated))))
            max_corrected_err=max(max_corrected_err,float(np.max(np.abs(vals-corrected))))
            boundary.append((s,float(np.max(vals))))
    # Continuous-time stability transition predicted at kE+kI=1, tested by sweep.
    stable_s=sorted(set(round(s,10) for s,m in boundary if m < -1e-8))
    unstable_s=sorted(set(round(s,10) for s,m in boundary if m > 1e-8))
    # Measure slow decay from a pure dominant eigenmode, avoiding mixtures.
    decay=[]; dt=.001
    for s in [.20,.40,.60,.75,.85,.92,.97]:
        kE,kI=s*.6,s*.4; J=jacobian(kE,kI)
        vals,vecs=np.linalg.eig(J); ix=np.argmax(vals.real)
        x=vecs[:,ix].real; x=x/np.linalg.norm(x); norms=[]
        for _ in range(3000):
            norms.append(np.linalg.norm(x)); x=x+dt*J@x
        slope=np.polyfit(np.arange(500,2500)*dt,np.log(norms[500:2500]),1)[0]
        decay.append([s,-float(slope),2*(1-s)])
    # Fixed point scales as (1-s)^-1 when D has projection on the critical mode.
    fp=[]; D=np.array([.002,.001,0.])
    for s in [.50,.60,.70,.78,.84,.88,.91,.93,.95,.96]:
        kE,kI=.6*s,.4*s; c=fixed_point(kE,kI,D); fp.append([1-s,float(np.linalg.norm(c))])
    slope=float(np.polyfit(np.log([q[0] for q in fp]),np.log([q[1] for q in fp]),1)[0])
    return {'stated_delta_eigenvalue_max_abs_error':max_stated_err,
      'corrected_sum_eigenvalue_max_abs_error':max_corrected_err,
      'claimed_boundary_delta':1.0,'actual_boundary_sum_kE_plus_kI':1.0,
      'stable_sweep_max_sum':max(stable_s),'unstable_sweep_min_sum':min(unstable_s),
      'decay_rate_rows_sum_measured_predicted':decay,
      'fixed_point_rows_margin_norm':fp,
      'fixed_point_loglog_slope_measured_predicted':[slope,-1.0]}

class TinyRNN(nn.Module):
    def __init__(self, vocab=8, h=24):
        super().__init__(); self.h=h; self.emb=nn.Embedding(vocab,h)
        self.W=nn.Parameter(torch.empty(h,h)); self.b=nn.Parameter(torch.zeros(h)); self.out=nn.Linear(h,vocab)
        nn.init.normal_(self.W,0,.55); nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias)
    def forward(self,x, controller=False, alpha=3., eps=.05, tau=.8):
        B,T=x.shape; h=torch.zeros(B,self.h,device=x.device); hs=[]
        for t in range(T):
            z=self.emb(x[:,t])+h@self.W.T+self.b; h=torch.tanh(z); hs.append(h)
        H=torch.stack(hs,1); logits=self.out(H)
        d=self.h//2
        # block spectral gains; singular values are differentiable and inexpensive at this size
        kE=torch.linalg.matrix_norm(self.W[:d,:d],ord=2); kI=torch.linalg.matrix_norm(self.W[d:,d:],ord=2)
        delta=kE-kI
        cov=H.reshape(-1,self.h)-H.reshape(-1,self.h).mean(0,keepdim=True)
        tr=(cov.square().sum()/max(cov.shape[0]-1,1))
        penalty=alpha*torch.relu(delta-(1-eps))**2 + .03*torch.log1p(tr/tau) if controller else 0.*tr
        return logits,H,delta,tr,penalty

def train_variant(controller, steps=220):
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        torch.manual_seed(SEED); model=TinyRNN().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        # deterministic synthetic Markov sequence: target is next symbol, long enough to expose state growth
        g=torch.Generator(device=device); g.manual_seed(SEED)
        losses=[]; traces=[]; deltas=[]
        for step in range(steps):
            B,T=64,32; x=torch.randint(0,8,(B,T),generator=g,device=device); y=(x+1)%8
            logits,H,d,tr,p=model(x,controller); loss=nn.functional.cross_entropy(logits.reshape(-1,8),y.reshape(-1))+p
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
            if step%20==0 or step==steps-1:
                losses.append(float(loss.detach().cpu())); traces.append(float(tr.detach().cpu())); deltas.append(float(d.detach().cpu()))
        # long horizon evaluation without update
        x=torch.randint(0,8,(64,128),generator=g,device=device); y=(x+1)%8
        with torch.no_grad():
            logits,H,d,tr,p=model(x,controller)
            longloss=float(nn.functional.cross_entropy(logits.reshape(-1,8),y.reshape(-1)).cpu())
            longtr=float(tr.cpu()); longd=float(d.cpu())
        return {'train_loss_checkpoints':losses,'trace_checkpoints':traces,'delta_checkpoints':deltas,
                'long_loss':longloss,'long_trace':longtr,'long_delta':longd,'device':device}
    except Exception as e:
        if device=='cuda':
            torch.cuda.empty_cache(); torch.manual_seed(SEED)
            # retry on CPU by temporarily disabling CUDA availability is awkward; call an explicit CPU helper
            old=torch.cuda.is_available
            torch.cuda.is_available=lambda: False
            try: return train_variant(controller,steps)
            finally: torch.cuda.is_available=old
        raise

def main():
    out={'toy_checks':toy_checks(),'baseline':train_variant(False),'controller':train_variant(True)}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()