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

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


def grad_loss(x):
    # Nonlinear scalar training dynamics: L(x)=.5*x^2 + .08*x^4.
    return x + .32*x**3


def jac_grad(x):
    return 1.0 + .96*x*x


def rollout(x0, eta, H):
    xs = [float(x0)]
    x = float(x0)
    for _ in range(H):
        x = x - eta * grad_loss(x)
        xs.append(x)
    return np.asarray(xs)


def rollout_sensitivity(x0, eta, H):
    # Exact tangent map R(t)=d F_t(x0)/d x0.
    r = 1.0
    rs = [r]
    x = float(x0)
    for _ in range(H):
        r *= 1.0 - eta * jac_grad(x)
        x = x - eta * grad_loss(x)
        rs.append(r)
    return np.asarray(rs)


def violation(gamma, xstar, eta, H, r=None):
    nom = rollout(xstar, eta, H)
    sens = rollout_sensitivity(xstar, eta, H)
    actual = rollout(xstar + gamma, eta, H)
    err = np.abs(actual - (nom + sens*gamma))
    return float(np.max(err)), float(err[-1])


def radial_boundary(xstar, eta, H, eps, max_r=2.0):
    # Exact one-dimensional analogue of ray bisection in the paper.
    lo, hi = 0.0, max_r
    if violation(hi, xstar, eta, H)[0] <= eps: return hi
    for _ in range(48):
        mid = (lo+hi)/2
        if violation(mid, xstar, eta, H)[0] <= eps: lo = mid
        else: hi = mid
    return lo


def toy_verification():
    xstar, eta, eps = 0.8, 0.12, 0.01
    rows=[]
    # Prediction 1: locally, error is quadratic in radius (log-log slope 2).
    rs=np.geomspace(.002,.12,12)
    es=np.array([violation(r,xstar,eta,4)[0] for r in rs])
    slope=float(np.polyfit(np.log(rs), np.log(es), 1)[0])
    rows.append({'prediction':'local violation scales as r^2','observed_loglog_slope':slope,'expected':2.0})
    # Prediction 2: trusted boundary scales sqrt(epsilon).
    epses=np.array([.0025,.005,.01,.02,.04])
    bounds=np.array([radial_boundary(xstar,eta,4,e) for e in epses])
    slope_b=float(np.polyfit(np.log(epses),np.log(bounds),1)[0])
    rows.append({'prediction':'boundary radius scales as sqrt(epsilon)','observed_loglog_slope':slope_b,'expected':.5,
                 'eps_radius_pairs':[[float(a),float(b)] for a,b in zip(epses,bounds)]})
    # Prediction 3: longer nonlinear rollout accumulates larger error.
    hs=np.arange(1,9)
    eh=np.array([violation(.08,xstar,eta,int(h))[0] for h in hs])
    rows.append({'prediction':'H-step error increases with horizon','horizons':hs.tolist(),'errors':eh.tolist(),
                 'monotone':bool(np.all(np.diff(eh)>=-1e-12))})
    return rows


def mlp_compare():
    # Tiny deterministic 2D binary task; numpy keeps the experiment reproducible.
    rng=np.random.RandomState(SEED)
    X=np.vstack([rng.randn(96,2)+[-1.0,-1.0], rng.randn(96,2)+[1.0,1.0]])
    y=np.r_[np.zeros(96),np.ones(96)]
    d, m = 2, 8
    n=d*m+m+m+1
    def unpack(p):
        i=0; W1=p[i:i+d*m].reshape(d,m); i+=d*m; b1=p[i:i+m]; i+=m
        W2=p[i:i+m]; i+=m; return W1,b1,W2,p[i]
    def loss_grad(p):
        W1,b1,W2,b2=unpack(p); q=np.tanh(X@W1+b1); logits=q@W2+b2
        prob=1/(1+np.exp(-np.clip(logits,-30,30)))
        loss=-np.mean(y*np.log(prob+1e-8)+(1-y)*np.log(1-prob+1e-8))
        dl=(prob-y)/len(y); dW2=q.T@dl; db2=dl.sum(); dq=dl[:,None]*W2
        da=dq*(1-q*q); dW1=X.T@da; db1=da.sum(0)
        return float(loss),np.r_[dW1.ravel(),db1,dW2,db2]
    p=rng.randn(n)*.12; pb=p.copy(); pt=p.copy()
    eta=.35; steps=40; eps=.0008; radius=.8; backward_baseline=0; backward_idea=0
    lb=[]; lt=[]; accepted=[]
    # The box coordinates are gradient and a fixed random orthogonal direction.
    qdir=rng.randn(n); qdir-=qdir.dot(np.ones(n))*0 # deterministic random direction
    qdir/=np.linalg.norm(qdir)
    for _ in range(steps):
        l,g=loss_grad(pb); backward_baseline+=1; lb.append(l); pb-=eta*g
        l,g=loss_grad(pt); backward_idea+=1; lt.append(l)
        gn=np.linalg.norm(g)+1e-12; b1=-g/gn
        # Orthogonalize a random direction to the gradient.
        b2=qdir-b1*qdir.dot(b1); b2/=np.linalg.norm(b2)+1e-12
        B=np.column_stack([b1,b2])
        # One-step nonlinear dynamics F(p)=p-eta grad L(p), and affine Jacobian action.
        f0=pt-eta*g
        R=np.empty((n,2)); fd=1e-4
        for j in range(2):
            _,gj=loss_grad(pt+fd*B[:,j]); backward_idea+=1
            R[:,j]=B[:,j]-eta*(gj-g)/fd
        def points(r):
            return np.array([[sx*r,sy*r] for sx in (-1,1) for sy in (-1,1)]+[[0,r],[0,-r],[r,0],[-r,0]])
        # Zeroth-order inner approximation of the trusted box by shrinking its radius.
        lo,hi=0.,radius
        for _ in range(12):
            mid=(lo+hi)/2; viol=0.
            for gamma in points(mid):
                _,gj=loss_grad(pt+B@gamma); backward_idea+=1
                actual=pt+B@gamma-eta*gj
                pred=f0+R@gamma
                viol=max(viol,float(np.linalg.norm(actual-pred)))
            if viol<=eps: lo=mid
            else: hi=mid
        rt=lo; accepted.append(rt)
        # Quadratic model: estimate directional Hessian action and minimize over tested trusted points.
        gammas=np.vstack([np.zeros((1,2)),points(rt)])
        vals=[]; delta=.02
        _,g1=loss_grad(pt+delta*b1); _,g2=loss_grad(pt+delta*b2); backward_idea+=2
        H11=max(0.,float((g1-g).dot(b1)/delta)); H22=max(0.,float((g2-g).dot(b2)/delta))
        for gamma in gammas:
            vals.append(l+g.dot(B@gamma)+.5*(H11*gamma[0]**2+H22*gamma[1]**2))
        pt+=B@gammas[int(np.argmin(vals))]
    return {'steps':steps,'baseline_final_loss':lb[-1],'trusted_final_loss':lt[-1],
            'baseline_best_loss':min(lb),'trusted_best_loss':min(lt),
            'mean_trusted_radius':float(np.mean(accepted)),
            'backward_evals_baseline':backward_baseline,'backward_evals_trusted':backward_idea,
            'baseline_loss_curve':lb,'trusted_loss_curve':lt}

def main():
    out={'seed':SEED,'toy':toy_verification(),'mlp':mlp_compare()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()