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

SEED = 1102
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(8, torch.get_num_threads()))
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')


def toy_verification():
    # A two-node positive cycle K has rho(K)=sqrt(ab), and tr(K^2)=2ab.
    a, b = 1.8, 0.8
    K = np.array([[0., a], [b, 0.]])
    rho = math.sqrt(a*b)
    trace2 = np.trace(K @ K)
    # Prediction 1: trace exactly equals the sum of the two length-2 cycles.
    p1_err = abs(trace2 - 2*a*b)
    # Prediction 2: repeated positive feedback changes from decay to growth at gamma*rho=1.
    gammas = np.array([0.35, 0.55, 0.70, 0.74, 0.76, 0.90, 1.10])
    n = 30
    growth = []
    for g in gammas:
        x = np.array([1., 0.])
        for _ in range(n): x = g*K @ x
        growth.append(float(np.linalg.norm(x)))
    predicted_threshold = 1.0/rho
    # Fit log(norm) vs steps at representative stable/unstable settings.
    slopes = []
    for g in [0.7, 0.9]:
        x=np.array([1., 0.]); vals=[]
        for _ in range(1, 16):
            x=g*K@x; vals.append(np.log(np.linalg.norm(x)+1e-30))
        slopes.append(float(np.polyfit(np.arange(1,16), vals, 1)[0]))
    # Prediction 3: normalized r=2 closed-walk term scales exactly as tau^-2.
    taus=np.array([0.5, 1., 2., 4.])
    normalized = np.array([trace2/(t*t*2.) for t in taus])
    slope_tau=float(np.polyfit(np.log(taus), np.log(normalized), 1)[0])
    # A disconnected/feed-forward matrix has no 2-cycle and zero trace(K^2).
    dag=np.array([[0., 2.0],[0., 0.]])
    return {
        'device': str(device),
        'predictions': {
            'trace_identity': {'predicted': 'tr(K^2)=2ab', 'observed': trace2, 'expected': 2*a*b, 'abs_error': p1_err},
            'threshold': {'predicted_gamma': predicted_threshold, 'observed_rho': rho,
                          'gamma_sweep': gammas.tolist(), 'norm_after_30': growth,
                          'classification': ['decay' if x < 1 else 'growth' for x in growth],
                          'theory_log_slopes_at_gamma_0.7_0.9': [math.log(g*rho) for g in [0.7,0.9]],
                          'observed_log_slopes': slopes},
            'temperature_scaling': {'predicted_loglog_slope': -2.0, 'observed_loglog_slope': slope_tau,
                                    'taus': taus.tolist(), 'normalized_trace': normalized.tolist()},
            'cycle_vanishes_in_dag': {'cycle_trace': float(np.trace(dag@dag)), 'predicted': 0.0}
        }
    }

class IterCell(nn.Module):
    def __init__(self):
        super().__init__()
        self.A=nn.Parameter(torch.tensor([[0.15, 0.85],[0.55, 0.10]], dtype=torch.float32))
        self.B=nn.Parameter(torch.randn(2,1)*0.25)
        self.out=nn.Linear(2,1)
    def forward(self,x, return_jac=False):
        h=torch.zeros(x.shape[0],2,device=x.device)
        for _ in range(5): h=torch.tanh(h@self.A.T + x@self.B.T)
        return self.out(h)
    def cycle_penalty(self, x, tau=1.0):
        # Exact per-example 2x2 block Jacobian for this small scalar-block MVP.
        h=torch.zeros(x.shape[0],2,device=x.device)
        for _ in range(4): h=torch.tanh(h@self.A.T + x@self.B.T)
        z=h@self.A.T + x@self.B.T
        d=(1-torch.tanh(z)**2).detach() # intermittent/stop-gradient Jacobian estimate
        # J_ij = |d tanh(z_i)/d h_j|; batch average gives nonnegative block matrix.
        J=d.mean(0)[:,None]*self.A.abs()
        return torch.trace(J@J)/(tau*tau*2), float(torch.trace(J@J).detach().cpu())
    def spectral_proxy(self):
        # Same auxiliary cost class: penalize total squared Jacobian entries, not cycles.
        return (self.A*self.A).mean()

def train(kind, seed=SEED, steps=500):
    torch.manual_seed(seed); np.random.seed(seed)
    model=IterCell().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=0.015)
    x=torch.randn(96,1,device=device)
    y=2*x + 0.35*torch.sin(3*x)
    t0=time.time(); losses=[]; penalties=[]
    for step in range(steps):
        # fixed stream makes variants directly comparable
        xb=x[(step*13)%96:((step*13)%96)+32] if (step*13)%96+32<=96 else x[:32]
        yb=y[(step*13)%96:((step*13)%96)+32] if (step*13)%96+32<=96 else y[:32]
        pred=model(xb); loss=((pred-yb)**2).mean()
        if kind=='cycle': pen,_=model.cycle_penalty(xb); loss=loss+0.04*pen
        elif kind=='spectral': pen=model.spectral_proxy(); loss=loss+0.04*pen
        else: pen=torch.tensor(0.,device=device)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10.); opt.step()
        losses.append(float(loss.detach().cpu())); penalties.append(float(pen.detach().cpu()))
    with torch.no_grad():
        pred=model(x); mse=float(((pred-y)**2).mean().cpu())
        cp, raw=model.cycle_penalty(x)
        sp=float(model.spectral_proxy().cpu())
    return {'final_mse':mse,'final_loss':losses[-1],'mean_aux':float(np.mean(penalties[-100:])),
            'closed_walk_trace2':raw,'spectral_proxy':sp,'seconds':time.time()-t0}

def main():
    out=toy_verification()
    train_results={k:train(k) for k in ['none','spectral','cycle']}
    out['mini_experiment']=train_results
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()