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

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


def sigmoid(x):
    x = np.clip(x, -60.0, 60.0)
    return 1.0 / (1.0 + np.exp(-x))


def toy_checks():
    # Prediction 1: q(E)=1/2 exactly at E=E_star and transition is about 4*tau.
    E_star, tau = 2.0, 0.25
    qstar = float(sigmoid((E_star-E_star)/tau))
    e10 = E_star - tau*math.log(9.0)
    e90 = E_star + tau*math.log(9.0)
    transition_width = e90-e10

    # Prediction 2: in the active, unforced regime dE/dt=-2*c*E,
    # hence log(E) slope=-2c, independently of initial energy.
    decay = []
    for c in [0.2, 0.5, 1.0, 1.5]:
        dt=1e-3; n=10000; v=math.sqrt(2*8.0); logs=[]; ts=[]
        for k in range(n):
            E=.5*v*v
            if k % 10 == 0:
                logs.append(math.log(E)); ts.append(k*dt)
            q=sigmoid((E-E_star)/tau)
            v += -dt*q*c*v
        # fit only well-active portion, E > 3 E_star
        mask=np.array([math.exp(x)>3*E_star for x in logs])
        slope=np.polyfit(np.array(ts)[mask],np.array(logs)[mask],1)[0]
        decay.append({'c':c,'predicted_log_slope':-2*c,'observed_log_slope':float(slope),
                      'relative_error':float(abs(slope+2*c)/(2*c))})

    # Prediction 3: bounded adversarial forcing g=-G sign(v), active equilibrium
    # has E*=G^2/(2c^2), from G*sqrt(2E)=2cE.
    forced=[]; G=1.0
    for c in [0.25,0.5,1.0,2.0]:
        dt=2e-4; n=250000; v=0.05
        vals=[]
        # Choose a low threshold so every predicted equilibrium is well inside q≈1.
        forced_E_star, forced_tau = 0.01, 0.001
        for k in range(n):
            E=.5*v*v; q=sigmoid((E-forced_E_star)/forced_tau)
            g=-G*(1.0 if v >= 0 else -1.0)
            v += dt*(-g-q*c*v)
            if k >= n//2: vals.append(.5*v*v)
        observed=float(np.mean(vals[-50000:]))
        predicted=G*G/(2*c*c)
        forced.append({'c':c,'predicted_plateau_E':predicted,
                       'observed_plateau_E':observed,
                       'relative_error':abs(observed-predicted)/predicted})
    return {'gate_transition':{'E_star':E_star,'tau':tau,'predicted_q_at_Estar':.5,
                               'observed_q_at_Estar':qstar,
                               'predicted_width_10_to_90':4.3944491547*tau,
                               'observed_width_10_to_90':transition_width},
            'unforced_decay':decay,'forced_plateau':forced}


def train_benchmark():
    # sklearn digits is a tiny, reproducible MNIST-like classification task.
    import torch
    import torch.nn as nn
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split
    torch.manual_seed(SEED); np.random.seed(SEED)
    try:
        device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # Trigger a tiny allocation so CUDA initialization errors are caught.
        if device.type=='cuda': torch.empty(1,device=device)
    except Exception:
        device=torch.device('cpu')
    x,y=load_digits(return_X_y=True)
    x=x.astype('float32')/16.0
    xa,xb,ya,yb=train_test_split(x,y,test_size=.25,random_state=SEED,stratify=y)
    tx=torch.tensor(xa); ty=torch.tensor(ya,dtype=torch.long)
    vx=torch.tensor(xb); vy=torch.tensor(yb,dtype=torch.long)
    def make():
        return nn.Sequential(nn.Linear(64,96),nn.ReLU(),nn.Linear(96,64),nn.ReLU(),nn.Linear(64,10)).to(device)
    def run(mode, lr):
        torch.manual_seed(SEED)
        model=make(); lossfn=nn.CrossEntropyLoss()
        # Explicit velocity implements theta <- theta + v, matching the stated interface.
        vel=[torch.zeros_like(p,device=device) for p in model.parameters()]
        bs=128; losses=[]; energies=[]; gates=[]
        t0=time.time()
        for step in range(300):
            idx=((torch.arange(bs)+step*bs) % len(tx))
            out=model(tx[idx].to(device)); loss=lossfn(out,ty[idx].to(device))
            gs=torch.autograd.grad(loss,tuple(model.parameters()))
            E=.5*sum(float((v*v).sum().detach().cpu()) for v in vel)
            q=float(sigmoid((E-0.02)/0.01)) if mode=='gated' else 0.0
            clip_scale=1.0
            if mode=='clip':
                total_norm=torch.sqrt(sum((g*g).sum() for g in gs))
                clip_scale=min(1.0, 1.0/(float(total_norm.detach().cpu())+1e-12))
            with torch.no_grad():
                for j,(p,g) in enumerate(zip(model.parameters(),gs)):
                    gg=g*clip_scale
                    vel[j].mul_(0.9*(1.0-lr*1.0*q)).add_(gg,alpha=-lr)
                    p.add_(vel[j])
            losses.append(float(loss.detach().cpu())); energies.append(E); gates.append(q)
        with torch.no_grad():
            acc=float((model(vx.to(device)).argmax(1).cpu()==vy).float().mean())
        return {'final_train_loss':losses[-1],'eval_accuracy':acc,'max_energy':max(energies),
                'mean_last50_energy':float(np.mean(energies[-50:])),'mean_gate_last50':float(np.mean(gates[-50:])),
                'seconds':time.time()-t0,'device':str(device)}
    results={}
    for lr in [0.03,0.1,0.3]:
        for mode in ['momentum','clip','gated']:
            try: results[f'{mode}_lr{lr}']=run(mode,lr)
            except Exception as e:
                results[f'{mode}_lr{lr}']={'error':repr(e)}
    return results

if __name__=='__main__':
    out={'seed':SEED,'toy':toy_checks()}
    try: out['benchmark']=train_benchmark()
    except Exception as e: out['benchmark_error']=repr(e)
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))