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

SEED = 915
np.random.seed(SEED); random.seed(SEED)


def pi_run(lam, eta, kp, ki, steps=500, reset=True, x0=1.0):
    # Exact stated ordering: theta uses I_k; the sign event determines I_{k+1}.
    x = float(x0); I = 0.0; prev_g = None
    xs, losses, Is, resets = [], [], [], []
    for k in range(steps):
        g = lam*x
        event = prev_g is not None and g*prev_g < 0
        x = x - eta*(kp*g + ki*I)
        next_I = 0.0 if (reset and event) else I + g
        xs.append(x); losses.append(.5*lam*x*x); Is.append(next_I); resets.append(int(event and reset))
        I = next_I
        prev_g = g
    return np.asarray(xs), np.asarray(losses), np.asarray(Is), np.asarray(resets)


def rho(lam, eta, kp, ki):
    A = np.array([[1-eta*kp*lam, -eta*ki], [lam, 1.]])
    return max(abs(np.linalg.eigvals(A)))


def toy_checks():
    kp, ki, lam = 1.0, .4, 1.0
    predicted_eta = 4/(lam*(2*kp-ki))
    # Prediction 1: spectral radius reaches 1 at the Jury boundary.
    etas = np.linspace(.1, 3.0, 2901)
    rhos = np.array([rho(lam,e,kp,ki) for e in etas])
    crossing = etas[np.argmin(abs(rhos-1))]
    # Prediction 2: varying lambda scales the boundary inversely.
    lambdas = np.array([.5, 1., 2., 4.])
    boundaries = []
    for l in lambdas:
        es = np.linspace(.05, 8/l, 2401)
        rr = np.array([rho(l,e,kp,ki) for e in es])
        boundaries.append(es[np.argmin(abs(rr-1))])
    pred_bounds = 4/(lambdas*(2*kp-ki))
    # Prediction 3: reset fires on the first gradient sign reversal and zeros old memory.
    # Choose an intentionally oscillatory proportional step; inspect the first event.
    x, loss, I, reset = pi_run(1., 1.8, 1., .4, 30, True)
    event_idx = np.flatnonzero(reset)
    first = int(event_idx[0]) if len(event_idx) else -1
    before_I = float(I[first-1]) if first > 0 else float('nan')
    at_I = float(I[first]) if first >= 0 else float('nan')
    # Prediction 3 diagnostic: reset truncates memory at reversals, but does not
    # generally enlarge the linear no-reset stability region.
    sweep=[]
    for eta in np.arange(.4, 2.41, .2):
        row={'eta':float(eta)}
        for name,flag in [('reset',True),('no_reset',False)]:
            xx,ll,ii,rr=pi_run(1.,float(eta),kp,ki,200,flag)
            settle=-1
            for j in range(len(ll)):
                if np.all(ll[j:] < 1e-8): settle=j; break
            row[name]={'resets':int(rr.sum()),'max_abs_x':float(np.max(np.abs(xx))),
                       'final_loss':float(ll[-1]),'settling_step':settle}
        sweep.append(row)
    return {
      'predicted_eta_boundary': predicted_eta, 'observed_eta_rho1': crossing,
      'boundary_relative_error': abs(crossing-predicted_eta)/predicted_eta,
      'lambda_sweep': [{'lambda':float(l),'predicted_eta':float(p),'observed_eta':float(o),'rel_error':float(abs(o-p)/p)} for l,p,o in zip(lambdas,pred_bounds,boundaries)],
      'reset_first_event_step': first, 'integral_before_reset': before_I, 'integral_at_reset': at_I,
      'eta_sweep':sweep
    }


def mlp_experiment():
    # Small fixed synthetic two-moons-like dataset, avoiding external data downloads.
    try:
        import torch
        from torch import nn
        torch.manual_seed(SEED); np.random.seed(SEED)
        dev = 'cuda' if torch.cuda.is_available() else 'cpu'
        n=512
        t=np.linspace(0, math.pi, n//2)
        X=np.vstack([np.c_[np.cos(t),np.sin(t)], np.c_[1-np.cos(t),1-np.sin(t)-.35]])
        X += .08*np.random.randn(n,2)
        y=np.r_[np.zeros(n//2),np.ones(n//2)].astype(np.int64)
        perm=np.random.RandomState(SEED).permutation(n); X=X[perm]; y=y[perm]
        Xt=torch.tensor(X,dtype=torch.float32,device=dev); yt=torch.tensor(y,device=dev)
        def train(kind):
            torch.manual_seed(SEED+ (1 if kind=='pi' else 0))
            m=nn.Sequential(nn.Linear(2,24),nn.Tanh(),nn.Linear(24,2)).to(dev)
            lossfn=nn.CrossEntropyLoss(); prev=None; I=[torch.zeros_like(p) for p in m.parameters()]
            losses=[]; resets=0
            for step in range(300):
                m.zero_grad(set_to_none=True); loss=lossfn(m(Xt),yt); loss.backward()
                gs=[p.grad.detach().clone() for p in m.parameters()]
                if kind=='sgd':
                    with torch.no_grad():
                        for p,g in zip(m.parameters(),gs): p -= .08*g
                else:
                    dot=sum((g*q).sum() for g,q in zip(gs,prev)) if prev is not None else 1.
                    if prev is not None and dot.item()<0:
                        I=[torch.zeros_like(p) for p in m.parameters()]; resets+=1
                    else: I=[a+g for a,g in zip(I,gs)]
                    with torch.no_grad():
                        for p,g,a in zip(m.parameters(),gs,I): p -= .04*(g+.4*a)
                    prev=gs
                losses.append(float(loss.detach().cpu()))
            return {'final_loss':losses[-1],'final_accuracy':float((m(Xt).argmax(1)==yt).float().mean().cpu()),'resets':resets}
        return {'device':dev,'sgd':train('sgd'),'pi_reset':train('pi')}
    except Exception as e:
        return {'error':repr(e),'device':'cpu'}

if __name__ == '__main__':
    out={'toy':toy_checks(),'mlp':mlp_experiment()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))