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

SEED=2196
np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda':
        torch.cuda.empty_cache()
        # force a tiny allocation so CUDA errors are caught here
        torch.zeros(1, device=device)
except Exception:
    device=torch.device('cpu')


def adjacent_radius(z):
    return torch.linalg.vector_norm(z[1:]-z[:-1], dim=1).max().item()

def recurrence_features(z, min_gap=2, top_k=8, margin=1e-4):
    """Cheap upper-triangular recurrence proxy. Lifetime is distance margin
    between a close-return threshold and the pair's distance; detached."""
    with torch.no_grad():
        D=torch.cdist(z,z)
        n=z.shape[0]; R=adjacent_radius(z)+margin
        vals=[]
        # local-in-time close returns, excluding the diagonal
        for i in range(n):
            for j in range(i+min_gap,n):
                d=float(D[i,j])
                if d < R:
                    # proxy filtration lifetime: threshold minus birth
                    vals.append((R-d,i,j))
        vals.sort(reverse=True)
        return [(float(w),i,j) for w,i,j in vals[:top_k]]

def cycle_loss(z, features):
    out=z.new_zeros(())
    for w,i,j in features:
        path=((z[i+1:j+1]-z[i:j])**2).sum()
        shortcut=((z[j]-z[i])**2).sum()
        out=out+float(w)*(path+shortcut)
    return out

def boundary(edges):
    b={}
    for a,c,coef in edges:
        b[a]=b.get(a,0)+coef; b[c]=b.get(c,0)-coef
    return {k:v for k,v in b.items() if v}

def loop_edges(i,j):
    # c_Gamma(i,j)-[i,j], oriented i -> ... -> j then j -> i
    e=[(k,k+1,1) for k in range(i,j)]
    e.append((j,i,1))
    return e

class Encoder(nn.Module):
    def __init__(self):
        super().__init__(); self.gru=nn.GRU(1,12,batch_first=True); self.proj=nn.Linear(12,3); self.head=nn.Linear(3,2)
    def forward(self,x):
        h,_=self.gru(x); z=self.proj(h); return z,self.head(z[:,-1])

def make_data(N=192,n=24):
    rng=np.random.default_rng(SEED); X=[]; y=[]
    for k in range(N):
        cls=k%2; t=np.linspace(0,2*np.pi,n)
        f=1.0 if cls==0 else 2.0
        sig=np.sin(f*t+rng.normal(0,.3))+rng.normal(0,.10,n)
        X.append(sig[:,None]); y.append(cls)
    p=rng.permutation(N); split=int(.75*N)
    return (torch.tensor(np.array(X)[p[:split]],dtype=torch.float32),torch.tensor(np.array(y)[p[:split]],dtype=torch.long),torch.tensor(np.array(X)[p[split:]],dtype=torch.float32),torch.tensor(np.array(y)[p[split:]],dtype=torch.long))

def _train_on(reg, dev, epochs=35):
    torch.manual_seed(SEED + (1 if reg else 0)); np.random.seed(SEED)
    xa,ya,xb,yb=make_data(); model=Encoder().to(dev); opt=torch.optim.Adam(model.parameters(),lr=.012)
    xa,ya,xb,yb=[q.to(dev) for q in (xa,ya,xb,yb)]
    t0=time.perf_counter()
    for _ in range(epochs):
        model.train(); opt.zero_grad(); z,logits=model(xa); loss=nn.functional.cross_entropy(logits,ya)
        if reg:
            for q in range(min(32,xa.shape[0])):
                f=recurrence_features(z[q]); loss=loss+0.004*cycle_loss(z[q],f)
        loss.backward(); opt.step()
    model.eval()
    with torch.no_grad(): pred=model(xb)[1].argmax(1); acc=(pred==yb).float().mean().item()
    return acc, time.perf_counter()-t0

def train(reg, epochs=35):
    try:
        return _train_on(reg, device, epochs)
    except Exception as exc:
        print(json.dumps({'cuda_fallback':type(exc).__name__}))
        return _train_on(reg, torch.device('cpu'), epochs)

def main():
    # deterministic algebraic toy trajectory
    z=torch.tensor([[0.,0.],[1.,0.],[1.,1.],[0.,1.],[0.,0.]],dtype=torch.float32)
    maxadj=adjacent_radius(z)
    b=boundary(loop_edges(0,4))
    boundary_ok=(b=={})
    margins=np.array([-0.01,0.0,0.0001,0.25,1.0])
    D=torch.cdist(z,z).numpy()
    # Prediction 1: every adjacent path edge is admitted exactly when r>R_gamma.
    path_ok=[bool(all(D[k,k+1] < maxadj+float(m) for k in range(len(z)-1))) for m in margins]
    path_predicted=[bool(m>0) for m in margins]
    # Prediction 2: the first nonadjacent recurrence appears at
    # m*=min_{j>=i+2} D_ij - R_gamma (strict inequality).
    nonadj=[D[i,j] for i in range(len(z)) for j in range(i+2,len(z))]
    predicted_nonadj=float(min(nonadj)-maxadj)
    recurrence_counts=[int(sum(D[i,j] < maxadj+float(m) for i in range(len(z)) for j in range(i+2,len(z)))) for m in margins]
    observed_nonadj=float(margins[next(k for k,c in enumerate(recurrence_counts) if c>0)]) if any(recurrence_counts) else None
    # Persistence/scale prediction: alpha multiplication gives alpha loss.
    zz=torch.tensor([[0.,0.],[.7,.1],[1.,.8],[.2,1.1],[0.,0.]],dtype=torch.float32)
    fs=[(0.8,0,4),(0.4,1,3)]
    base=float(cycle_loss(zz,fs)); scales=[0.,.25,.5,1.,2.]
    losses=[float(cycle_loss(zz,[(a*w,i,j) for w,i,j in fs])) for a in scales]
    ratios=[(v/base if base else 0.) for v in losses]
    # A second prediction: zero selected features => zero regularizer.
    zero=float(cycle_loss(zz,[]))
    baseline=train(False); idea=train(True)
    result={'device':str(device),'math':{
      'boundary_identity_predicted_zero':0,'boundary_residual':b,'boundary_confirmed':boundary_ok,
      'R_gamma':maxadj,'path_admission_prediction':'all adjacent edges iff margin > 0',
      'path_admission_observed':path_ok,'path_admission_predicted':path_predicted,
      'nonadjacent_transition_prediction':predicted_nonadj,
      'margins':margins.tolist(),'recurrence_vertex_counts':recurrence_counts,
      'observed_nonadjacent_transition_margin':observed_nonadj,
      'scale_prediction':'L(alpha*w)/L(w)=alpha; L(empty)=0','scales':scales,'losses':losses,'ratios':ratios,
      'zero_feature_loss':zero},
      'mini_experiment':{'baseline_accuracy':baseline[0],'idea_accuracy':idea[0],'baseline_seconds':baseline[1],'idea_seconds':idea[1]}}
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()