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

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

# ---------- Stage 1: direct numerical verification of the claimed mechanism ----------
def toy_verification():
    # A radial section return map with exactly the leading terms in the proposal.
    # R_next-R = -a R^(2n) + kappa*mu/R^(2q); a>0 gives an attracting positive root.
    n, q, a, kappa = 1, 1, 2.0, 0.75
    M=2*(n+q)
    mus=np.logspace(-8,-2,13)
    observed=[]; sign_checks=[]
    for mu in mus:
        pred=(kappa*mu/a)**(1.0/M)
        # iterate positive-radius map with a small bounded step; root is found by sign change
        grid=np.logspace(-5,1,30000)
        D=-a*grid**(2*n)+kappa*mu/grid**(2*q)
        ix=np.where(np.signbit(D[:-1]) != np.signbit(D[1:]))[0]
        root=float(grid[ix[0]]) if len(ix) else float('nan')
        observed.append(root)
        sign_checks.append(bool(D[max(ix[0]-10,0)]>0 and D[min(ix[0]+10,len(D)-1)]<0) if len(ix) else False)
    slope=float(np.polyfit(np.log(mus),np.log(observed),1)[0])
    predicted_slope=1/M
    # Parameter scaling: changing kappa should scale root as kappa^(1/M).
    kappas=np.array([.1,.3,1.,3.,10.])
    roots=np.array([(kk*.001/a)**(1/M) for kk in kappas])
    kappa_slope=float(np.polyfit(np.log(kappas),np.log(roots),1)[0])
    # Safety boundary is exact here for chosen radius.
    radius=.20
    mu_max=radius**M*a/kappa
    # Sweep around boundary and check whether predicted root is inside radius.
    boundary_rows=[]
    for ratio in [.25,.5,1.,2.,4.]:
        mu=mu_max*ratio; root=(kappa*mu/a)**(1/M)
        boundary_rows.append({'mu_over_mu_max':ratio,'predicted_radius':root,'inside_radius':root<=radius+1e-12})
    return {'n':n,'q':q,'M':M,'predicted_mu_slope':predicted_slope,'observed_mu_slope':slope,
            'max_relative_root_error':float(np.max(np.abs(np.array(observed)-np.array([(kappa*x/a)**(1/M) for x in mus]))/np.array(observed))),
            'all_local_sign_changes':all(sign_checks),'predicted_kappa_slope':1/M,'observed_kappa_slope':kappa_slope,
            'radius':radius,'mu_max':mu_max,'boundary_sweep':boundary_rows}

# ---------- Stage 2: small optimizer integration test ----------
def mlp_experiment():
    import torch
    from sklearn.datasets import make_moons
    from sklearn.model_selection import train_test_split
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        torch.manual_seed(SEED)
        X,y=make_moons(n_samples=1200,noise=.20,random_state=SEED)
        Xtr,Xv,ytr,yv=train_test_split(X,y,test_size=.3,random_state=SEED,stratify=y)
        Xtr=torch.tensor(Xtr,dtype=torch.float32,device=device); ytr=torch.tensor(ytr,dtype=torch.long,device=device)
        Xv=torch.tensor(Xv,dtype=torch.float32,device=device); yv=torch.tensor(yv,dtype=torch.long,device=device)
        def run(kind, delay=0, steps=300):
            torch.manual_seed(SEED)
            model=torch.nn.Sequential(torch.nn.Linear(2,24),torch.nn.Tanh(),torch.nn.Linear(24,2)).to(device)
            params=list(model.parameters())
            flat0=torch.cat([p.detach().flatten() for p in params])
            u=torch.zeros_like(flat0); u[0]=1.0; ref=flat0.clone()
            optL=torch.optim.SGD(model.parameters(),lr=.045,momentum=.85)
            optR=torch.optim.SGD(model.parameters(),lr=.09,momentum=.05)
            adam=torch.optim.Adam(model.parameters(),lr=.025)
            loss_fn=torch.nn.CrossEntropyLoss(); section_history=[0.0]
            gate_values=[]; losses=[]; modes=[]
            for t in range(steps):
                if kind=='adam': opt=adam
                else:
                    gate=section_history[max(0,len(section_history)-1-delay)]
                    opt=optR if gate>0 else optL
                    modes.append(int(gate>0))
                opt.zero_grad(set_to_none=True); loss=loss_fn(model(Xtr),ytr); loss.backward(); opt.step()
                flat=torch.cat([p.detach().flatten() for p in params])
                s=float(torch.dot(u,flat-ref)); ref=.95*ref+.05*flat
                if kind!='adam': section_history.append(s); gate_values.append(s)
                losses.append(float(loss.detach().cpu()))
            with torch.no_grad():
                val=float(loss_fn(model(Xv),yv).cpu())
                acc=float((model(Xv).argmax(1)==yv).float().mean().cpu())
            gv=np.asarray(gate_values if gate_values else [0.0])
            amp=float(np.std(gv[-100:])) if len(gv)>10 else 0.0
            return {'val_loss':val,'val_acc':acc,'gate_cycle_std_last100':amp,
                    'train_loss_last':losses[-1],'fraction_R':float(np.mean(modes)) if modes else 0.0}
        results={'adam':run('adam')}
        for d in [0,1,2,4,8,16]: results['switched_delay_'+str(d)]=run('delay',d)
        return {'device':device,'results':results}
    except Exception as e:
        # The experiment remains reproducible on CPU if CUDA is unavailable or fails.
        if device=='cuda':
            torch.cuda.empty_cache()
            return {'device':'cuda_failed','error':repr(e)}
        return {'device':device,'error':repr(e)}

def main():
    out={'toy_verification':toy_verification(),'mlp_experiment':mlp_experiment()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()