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

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


def finite_horizon_sweep():
    # z_{k+1}=a z_k, V=.5 z^2. Theory: V_{k+M}/V_k=a^(2M),
    # and alpha-contraction boundary |a|=(1-alpha)^(1/(2M)).
    alpha=0.10
    rows=[]
    for M in (2,4,8):
        pred=(1-alpha)**(1/(2*M))
        aa=np.linspace(0.80,1.06,261)
        ratios=aa**(2*M)
        ok=ratios <= 1-alpha
        obs=aa[np.where(ok)[0][-1]]
        rows.append({'M':M,'predicted_boundary':pred,'observed_grid_boundary':float(obs),
                     'abs_error':float(abs(obs-pred))})
    # Scaling prediction: log(V_M/V_0)=2M log|a|.
    scale=[]
    for a in (0.82,0.90,0.98):
        vals=[]
        for M in (1,2,4,8):
            vals.append((M, float(a**(2*M)), float(2*M*math.log(a))))
        scale.append({'a':a,'values':vals})
    return rows, scale


def stochastic_mismatch_sweep():
    # For a=0, z_{k+1}=noise, residual D=V_{k+M}-V_k+alpha V_k.
    # E|D| scales quadratically with noise amplitude, as V is quadratic.
    alpha=.1; M=4; n=12000; burn=200
    out=[]
    for sigma in (.02,.05,.10,.20):
        rng=np.random.default_rng(100+int(sigma*1000))
        z=rng.normal(0,sigma,size=n+M+1)
        V=.5*z*z
        D=V[M:]-V[:-M]+alpha*V[:-M]
        d=D[burn:]
        ema=0.; beta=.9; eps=[]
        for x in np.abs(d):
            ema=beta*ema+(1-beta)*x; eps.append(ema)
        out.append({'sigma':sigma,'mean_abs_residual':float(np.mean(np.abs(d))),
                    'mean_ema_last_half':float(np.mean(eps[len(eps)//2:])),
                    'normalized_by_sigma2':float(np.mean(np.abs(d))/sigma**2)})
    return out


def unstable_ema_sweep():
    # Prediction: the EMA allowance is nearly harmless for stationary residuals,
    # but unstable exponential growth outruns it and creates violations.
    alpha=.1; M=4; n=120; beta=.9; out=[]
    for a in (.90,.98,1.00,1.01,1.03,1.06):
        z=0.02; V=[.5*z*z]
        for k in range(n+M): z=a*z; V.append(.5*z*z)
        D=np.array(V[M:])-np.array(V[:-M])+alpha*np.array(V[:-M])
        ema=0.; viol=[]; eps=[]
        for d in D:
            ema=beta*ema+(1-beta)*abs(d)
            viol.append(d-ema>0); eps.append(ema)
        out.append({'a':a,'theory_ratio':a**(2*M),
                    'violation_rate_last_half':float(np.mean(viol[len(viol)//2:])),
                    'last_residual_over_ema':float(D[-1]/max(eps[-1],1e-30))})
    return out


def mlp_experiment():
    # Tiny digits MLP. The idea is implemented as an online delayed gradient
    # Lyapunov penalty; current gradient is differentiable, old gradient and EMA
    # allowance are detached. This is a practical proxy for the optimizer test.
    import torch
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split
    torch.manual_seed(SEED); np.random.seed(SEED)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        X,y=load_digits(return_X_y=True)
        X=X.astype('float32')/16.; y=y.astype('int64')
        xt,xv,yt,yv=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
        def run(use_reg):
            torch.manual_seed(SEED)
            model=torch.nn.Sequential(torch.nn.Linear(64,64),torch.nn.Tanh(),torch.nn.Linear(64,10)).to(device)
            opt=torch.optim.SGD(model.parameters(),lr=.35,momentum=.0)
            lossfn=torch.nn.CrossEntropyLoss(); M=4; alpha=.1; lam=.03; beta=.9
            qs=[]; eps=0.; losses=[]; gnorms=[]; violations=[]
            order=np.arange(len(xt)); rng=np.random.default_rng(SEED)
            for epoch in range(12):
                rng.shuffle(order)
                for start in range(0,len(order),64):
                    ids=order[start:start+64]
                    xb=torch.tensor(xt[ids],device=device); yb=torch.tensor(yt[ids],device=device)
                    opt.zero_grad(set_to_none=True)
                    loss=lossfn(model(xb),yb)
                    grads=torch.autograd.grad(loss,tuple(model.parameters()),create_graph=use_reg,retain_graph=True)
                    flat=torch.cat([g.reshape(-1) for g in grads])
                    V=.5*(flat*flat).sum()
                    q_det=flat.detach()
                    penalty=torch.zeros((),device=device)
                    if use_reg and len(qs)>=M:
                        old=qs[-M]
                        Vold=.5*(old*old).sum()
                        D=V-Vold+alpha*Vold
                        eps=beta*eps+(1-beta)*float(abs(D.detach()).cpu())
                        r=D-eps
                        penalty=lam*torch.relu(r).clamp(max=10.)**2/(1.+Vold)
                        violations.append(float((r.detach()>0).cpu()))
                    total=loss+penalty
                    total.backward(); opt.step()
                    qs.append(q_det)
                    losses.append(float(loss.detach().cpu())); gnorms.append(float(torch.linalg.vector_norm(flat.detach()).cpu()))
            with torch.no_grad():
                pred=model(torch.tensor(xv,device=device)).argmax(1).cpu().numpy()
            acc=float(np.mean(pred==yv)); tail=np.array(losses[-100:]); gn=np.array(gnorms[-100:])
            return {'accuracy':acc,'final_loss':float(np.mean(tail)),
                    'loss_std_tail':float(np.std(tail)),'grad_std_tail':float(np.std(gn)),
                    'violation_rate':float(np.mean(violations)) if violations else 0.0}
        try:
            b=run(False); r=run(True)
        except Exception:
            if device=='cuda':
                torch.cuda.empty_cache(); device='cpu'; b=run(False); r=run(True)
            else: raise
        return {'device':device,'baseline':b,'idea':r}
    except Exception as e:
        return {'error':repr(e)}

if __name__=='__main__':
    result={'finite_horizon_boundary':finite_horizon_sweep()[0],
            'ratio_scaling':finite_horizon_sweep()[1],
            'stochastic_mismatch':stochastic_mismatch_sweep(),
            'unstable_ema':unstable_ema_sweep(),
            'mlp':mlp_experiment()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))