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

SEED = 7
torch.manual_seed(SEED); np.random.seed(SEED)
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')

# A small input-conditioned linear recurrent layer: h' = A(z)h + g(z).
class ICRec(nn.Module):
    def __init__(self, d=8, k=3, mode='free', rho=.90, lam=3.0):
        super().__init__(); self.d=d; self.k=k; self.mode=mode; self.rho=rho; self.lam=lam
        self.A = nn.Parameter(.25*torch.randn(k,d,d))
        self.inp = nn.Linear(1,d)
        self.gate = nn.Linear(1,k)
        self.out = nn.Linear(d,1)
        # P = L L^T + eps I, with positive diagonal through softplus.
        self.rawL = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
        self.log_diag = nn.Parameter(torch.zeros(d))

    def P(self):
        L = torch.tril(self.rawL, diagonal=-1) + torch.diag(torch.nn.functional.softplus(self.log_diag)+.08)
        return L @ L.T + 1e-4*torch.eye(self.d, device=L.device)

    def matrices(self):
        if self.mode == 'spectral':
            # Standard pointwise Euclidean spectral normalization, applied separately.
            vals=[]
            for a in self.A:
                s=torch.linalg.svdvals(a)[0]
                vals.append(a * (.90 / (s + 1e-6)).clamp(max=1.0))
            return torch.stack(vals)
        return self.A

    def lyap_penalty(self):
        P=self.P(); mats=self.matrices()
        # Generalized eigenvalues of A^T P A v = lambda P v.
        # Cholesky transform is symmetric and differentiable for this small MVP.
        L=torch.linalg.cholesky(P)
        invL=torch.linalg.inv(L)
        vals=[]
        for a in mats:
            B=invL @ (a.T @ P @ a) @ invL.T
            vals.append(torch.linalg.eigvalsh((B+B.T)/2)[-1])
        vmax=torch.stack(vals)
        violation=torch.relu(vmax-self.rho**2)
        return self.lam*(violation**2).sum(), vmax.detach(), violation.detach()

    def forward(self,x, return_stats=False):
        h=torch.zeros(x.shape[0],self.d,device=x.device); hs=[]
        mats=self.matrices()
        for t in range(x.shape[1]):
            z=x[:,t,:]
            pi=torch.softmax(self.gate(z),-1)
            A=torch.einsum('bk,kij->bij',pi,mats)
            h=torch.bmm(A,h.unsqueeze(-1)).squeeze(-1)+self.inp(z)
            h=torch.tanh(h)
            hs.append(h)
        y=self.out(h).squeeze(-1)
        if return_stats: return y, torch.stack(hs,1)
        return y

def math_check():
    # Random positive P and matrices deliberately rescaled to satisfy the CQLF bound.
    torch.manual_seed(SEED+1); d=5; k=4; rho=.82
    L=torch.randn(d,d); P=L@L.T+.5*torch.eye(d); C=torch.linalg.cholesky(P)
    mats=[]
    for _ in range(k):
        a=torch.randn(d,d)
        # scale in P geometry so max generalized eigenvalue is below rho^2.
        invC=torch.linalg.inv(C)
        B=invC @ (a.T@P@a) @ invC.T
        lm=torch.linalg.eigvalsh((B+B.T)/2)[-1]
        mats.append(a*(rho*.72/torch.sqrt(lm)))
    mats=torch.stack(mats)
    # Every arbitrary convex gate and every length-25 product must contract in P norm.
    max_violation=-1.; max_ratio=0.; worst=0.
    for trial in range(200):
        h=torch.randn(d); e0=h@P@h
        for t in range(25):
            pi=torch.softmax(torch.randn(k),0); a=torch.einsum("k,kij->ij", pi, mats); h=a@h
            ratio=torch.sqrt((h@P@h)/e0).item()
            max_ratio=max(max_ratio, ratio/(rho**(t+1)))
        # ratio/rho^T <= 1 is the claimed bound
    for a in mats:
        invC=torch.linalg.inv(C); B=invC@(a.T@P@a)@invC.T
        max_violation=max(max_violation,(torch.linalg.eigvalsh((B+B.T)/2)[-1]-rho**2).item())
    return {'max_basis_violation':max_violation, 'max_normalized_product_ratio':max_ratio,
            'rho':rho, 'pass': bool(max_violation <= 2e-5 and max_ratio <= 1.0001)}

def train(mode, steps=350):
    torch.manual_seed(SEED); np.random.seed(SEED)
    m=ICRec(mode=mode).to(device); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
    t0=time.time(); losses=[]; final_h=0.; final_grad=0.; penalty=0.
    for step in range(steps):
        # First symbol is the bit to remember; remaining symbols are zero.
        x=torch.zeros(96,30,1,device=device); bit=torch.randint(0,2,(96,),device=device).float()
        x[:,0,0]=bit
        target=2*bit-1
        pred,h=m(x,True); task=nn.functional.mse_loss(pred,target)
        p,_,_=m.lyap_penalty() if mode=='lyap' else (torch.tensor(0.,device=device),None,None)
        loss=task+p; opt.zero_grad(); loss.backward()
        final_grad=float(torch.nn.utils.clip_grad_norm_(m.parameters(),100.).detach().cpu()); opt.step()
        losses.append(float(task.detach().cpu())); final_h=float(h.detach().abs().max().cpu()); penalty=float(p.detach().cpu())
    # Diagnostic random long products using learned basis matrices.
    with torch.no_grad():
        mats=m.matrices(); P=m.P(); C=torch.linalg.cholesky(P); invC=torch.linalg.inv(C)
        eig=[]
        for a in mats:
            B=invC@(a.T@P@a)@invC.T; eig.append(float(torch.linalg.eigvalsh((B+B.T)/2)[-1].cpu())**.5)
        x=torch.zeros(256,60,1,device=device); bit=torch.randint(0,2,(256,),device=device).float(); x[:,0,0]=bit
        pred,h=m(x,True); acc=float(((pred>0)==(bit>0)).float().mean().cpu())
    return {'loss_start':losses[0], 'loss_end':losses[-1], 'accuracy_long':acc,
            'max_hidden_last_batch':final_h, 'last_grad_norm':final_grad,
            'lyap_penalty':penalty, 'max_basis_P_gain':max(eig), 'seconds':time.time()-t0}

def main():
    check=math_check(); results={}
    for mode in ['free','spectral','lyap']:
        results[mode]=train(mode)
    out={'device':str(device), 'math_check':check, 'results':results}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()