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

SEED=274
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if DEVICE=='cuda': torch.cuda.get_device_properties(0)
except Exception:
    DEVICE='cpu'

# -------- Core finite-dimensional quotient verification --------
def quotient(A,u,v,eps=1e-12):
    return float(v @ ((np.eye(A.shape[0])-A)@u) / (v@u+eps))

def verify_math():
    rng=np.random.default_rng(19)
    # A positive, non-symmetric matrix has positive right/left Perron vectors.
    A=rng.uniform(.02,.9,(8,8)); np.fill_diagonal(A, .35)
    eig, vr=np.linalg.eig(A); j=np.argmax(eig.real)
    mu=float(eig[j].real); u=np.abs(vr[:,j].real); u/=np.linalg.norm(u)
    eigT, vl=np.linalg.eig(A.T); k=np.argmin(np.abs(eigT-mu))
    v=np.abs(vl[:,k].real); v/=np.linalg.norm(v)
    q=quotient(A,u,v); rr=np.linalg.norm((np.eye(8)-A)@u-(1-mu)*u)
    rl=np.linalg.norm((np.eye(8)-A.T)@v-(1-mu)*v)
    # Independent positive probes optimized by alternating gradient descent on residuals.
    U=np.exp(rng.normal(size=8)); V=np.exp(rng.normal(size=8)); U/=np.linalg.norm(U); V/=np.linalg.norm(V)
    before=(np.linalg.norm((np.eye(8)-A)@U-quotient(A,U,V)*U)**2+
            np.linalg.norm((np.eye(8)-A.T)@V-quotient(A,U,V)*V)**2)
    for _ in range(300):
        lam=quotient(A,U,V)
        # projected gradient-like updates for the two residuals, followed by positivity.
        r=(np.eye(8)-A)@U-lam*U
        s=(np.eye(8)-A.T)@V-lam*V
        U=np.maximum(U-0.08*((np.eye(8)-A).T@r-lam*r),1e-8); U/=np.linalg.norm(U)
        V=np.maximum(V-0.08*((np.eye(8)-A.T).T@s-lam*s),1e-8); V/=np.linalg.norm(V)
    after=(np.linalg.norm((np.eye(8)-A)@U-quotient(A,U,V)*U)**2+
           np.linalg.norm((np.eye(8)-A.T)@V-quotient(A,U,V)*V)**2)
    return {'positive_matrix_mu':mu,'quotient_at_eigenpair':q,
            'quotient_abs_error':abs(q-(1-mu)), 'right_residual':rr,
            'left_residual':rl,'probe_residual_before':before,
            'probe_residual_after':after,'probe_residual_ratio':after/before}

class RNN(nn.Module):
    def __init__(self,h=32):
        super().__init__(); self.W=nn.Parameter(torch.randn(h,h)*.18)
        self.U=nn.Parameter(torch.randn(h,2)*.25); self.b=nn.Parameter(torch.zeros(h))
        self.out=nn.Linear(h,1)
    def forward(self,x):
        h=torch.zeros(x.shape[0],self.W.shape[0],device=x.device)
        for t in range(x.shape[1]): h=torch.tanh(h@self.W.T+x[:,t]@self.U.T+self.b)
        return self.out(h).squeeze(-1)

def cone_penalty(model, steps=5, mu_max=.98):
    # Local Jacobian at the zero-input fixed point, which is the recurrent
    # transition relevant during the long distractor portion of this task.
    W=model.W; h=W.shape[0]
    with torch.no_grad(): A=W.detach().clone()
    a=torch.zeros(h,device=W.device,requires_grad=True); b=torch.zeros(h,device=W.device,requires_grad=True)
    opt=torch.optim.SGD([a,b],lr=.35)
    I=torch.eye(h,device=W.device)
    for _ in range(steps):
        u=F.softplus(a)+1e-5; v=F.softplus(b)+1e-5
        u=u/(u.norm()+1e-12); v=v/(v.norm()+1e-12)
        lam=((u-A@u)*v).sum()/(u@v+1e-6)
        loss=((u-A@u-lam*u)**2).sum()+((v-A.T@v-lam*v)**2).sum()
        opt.zero_grad(); loss.backward(); opt.step()
    u=(F.softplus(a).detach()+1e-5); u=u/u.norm()
    v=(F.softplus(b).detach()+1e-5); v=v/v.norm()
    # Reattach only the model matrix for the outer update.
    lam=((u-W@u)*v).sum()/(u@v+1e-6)
    mu=1-lam
    return F.softplus(mu-mu_max)**2, float(mu.detach()), float(lam.detach())

def spectral_penalty(model, target=.98):
    # Standard differentiable power estimate of the recurrent spectral norm.
    s=torch.linalg.matrix_norm(model.W,ord=2)
    return F.softplus(s-target)**2, float(s.detach())

def make_batch(n=64,T=30,device='cpu'):
    x=torch.zeros(n,T,2,device=device); y=torch.empty(n,device=device).uniform_(-1,1)
    x[:,0,0]=1.; x[:,0,1]=y
    return x,y

def train(kind,seed,steps=350):
    torch.manual_seed(seed); model=RNN().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    losses=[]; penalties=[]; t0=time.time()
    for step in range(steps):
        x,y=make_batch(device=DEVICE); pred=model(x); task=F.mse_loss(pred,y); reg=pred.new_zeros(())
        if kind=='cone': reg,mu,_=cone_penalty(model); penalties.append(mu)
        elif kind=='spectral': reg,s=spectral_penalty(model); penalties.append(s)
        loss=task+(.12 if kind!='baseline' else 0.)*reg
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
        losses.append(float(task.detach()))
    # long-horizon heldout loss and gradient statistic, same fixed batch.
    torch.manual_seed(1000+seed); x,y=make_batch(256,30,DEVICE)
    with torch.no_grad(): test=float(F.mse_loss(model(x),y))
    return {'final_train_loss':float(np.mean(losses[-30:])),'test_loss':test,
            'penalty_metric':float(np.mean(penalties[-30:])) if penalties else None,
            'seconds':time.time()-t0}

def main():
    mathcheck=verify_math(); results={'device':DEVICE,'math_check':mathcheck,'runs':{}}
    for kind in ['baseline','spectral','cone']:
        results['runs'][kind]=[train(kind,s) for s in [7,17,27]]
    for kind,rs in list(results['runs'].items()):
        if kind.endswith('_summary'): continue
        results['runs'][kind+'_summary']={k:float(np.mean([r[k] for r in rs])) for k in ['final_train_loss','test_loss','seconds']}
        if kind!='baseline': results['runs'][kind+'_summary']['penalty_metric']=float(np.mean([r['penalty_metric'] for r in rs]))
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))
if __name__=='__main__': main()