import json
import numpy as np
from scipy.optimize import linprog

SEED=11
rng=np.random.default_rng(SEED)
mu,sigma=0.055,0.20
U_EDGE,T_EDGE=-0.80,0.80
x=np.linspace(-1.,1.,201); N=len(x)
U=x<=U_EDGE; T=x>=T_EDGE; C=~(U|T); X0=(x>=-.2)&(x<=.2)
z,qw=np.polynomial.hermite.hermgauss(61); z*=np.sqrt(2.); qw/=np.sqrt(np.pi)
uix,tix=np.flatnonzero(U)[0],np.flatnonzero(T)[-1]
P=np.zeros((N,N))
for i,xi in enumerate(x):
    if U[i]: P[i,uix]=1.; continue
    if T[i]: P[i,tix]=1.; continue
    for y,wt in zip(xi+mu+sigma*z,qw):
        if y<=U_EDGE: P[i,uix]+=wt
        elif y>=T_EDGE: P[i,tix]+=wt
        else:
            j=np.clip(np.searchsorted(x,y)-1,0,N-2); a=(y-x[j])/(x[j+1]-x[j])
            P[i,j]+=wt*(1-a); P[i,j+1]+=wt*a

def oracle(delta):
    tv=N; c=np.zeros(N+1); c[tv]=1.; A=[]; b=[]
    for i in np.flatnonzero(C):
        r=np.zeros(N+1); r[:N]=P[i]; r[i]-=1.; A.append(r); b.append(-delta)
    for i in np.flatnonzero(X0):
        r=np.zeros(N+1); r[i]=1.; r[tv]=-1.; A.append(r); b.append(0.)
    bounds=[(1,1) if U[i] else ((0,0) if T[i] else (0,1)) for i in range(N)]+[(0,1)]
    r=linprog(c,A_ub=np.asarray(A),b_ub=np.asarray(b),bounds=bounds,method='highs')
    return (None,None) if not r.success else (float(r.x[tv]),r.x[:N])

def exact_failure():
    ci=np.flatnonzero(C); h=np.zeros(N); h[U]=1
    h[ci]=np.linalg.solve(np.eye(len(ci))-P[np.ix_(ci,ci)],P[np.ix_(ci,[uix])].ravel())
    return h

def margin_sweep():
    out=[]
    for d in [0,.002,.005,.01,.02,.03,.04,.05]:
        t,B=oracle(d); out.append({'delta':d,'min_initial_B':t,'p_cert':None if t is None else 1-t,'feasible':t is not None})
    return out

def mc_scaling():
    def bf(y): return np.clip((y+1)/2,0,1)
    xx=.1; truth=float(np.sum(qw*bf(xx+mu+sigma*z))); out=[]
    for M in [8,16,32,64,128,256,512,1024]:
        ee=[]
        for _ in range(400): ee.append(abs(bf(xx+mu+sigma*rng.normal(size=M)).mean()-truth))
        out.append({'M':M,'mae':float(np.mean(ee))})
    slope=float(np.polyfit(np.log([a['M'] for a in out]),np.log([a['mae'] for a in out]),1)[0])
    return truth,out,slope

def neural_comparison(h,delta=.02,steps=2500):
    try:
        import torch
        torch.manual_seed(SEED); torch.set_num_threads(4)
        dev='cuda' if torch.cuda.is_available() else 'cpu'
        # Small scalar MLP; P is a deterministic quadrature estimate of the stochastic expectation.
        xx=torch.tensor(x[:,None],dtype=torch.float32,device=dev); PP=torch.tensor(P,dtype=torch.float32,device=dev)
        UU=torch.tensor(U,device=dev); TT=torch.tensor(T,device=dev); CC=torch.tensor(C,device=dev); XX=torch.tensor(X0,device=dev)
        def run(kind):
            net=torch.nn.Sequential(torch.nn.Linear(1,32),torch.nn.Tanh(),torch.nn.Linear(32,32),torch.nn.Tanh(),torch.nn.Linear(32,1),torch.nn.Sigmoid()).to(dev)
            opt=torch.optim.Adam(net.parameters(),lr=3e-3)
            target=torch.tensor(h[:,None],dtype=torch.float32,device=dev)
            for k in range(steps):
                B=net(xx).squeeze(1); PB=PP@B
                boundary=((B[UU]-1)**2).mean()+(B[TT]**2).mean()
                if kind=='baseline': loss=boundary+2*((B-target)**2).mean()
                else:
                    sp=torch.nn.functional.softplus
                    loss=8*boundary+8*sp(B[XX]-(1-delta)).mean()+8*sp(PB[CC]-B[CC]+delta).mean()
                opt.zero_grad(); loss.backward(); opt.step()
            with torch.no_grad():
                B=net(xx).squeeze(1); PB=PP@B
                return {'loss':float(loss),'max_init_B':float(B[XX].max()),'max_init_limit_violation':float((B[XX]-(1-delta)).max()),'max_drift_residual':float((PB[CC]-B[CC]+delta).max()),'boundary_rmse':float(torch.sqrt(((B[UU]-1)**2).mean()+(B[TT]**2).mean()))}
        return {'device':dev,'baseline':run('baseline'),'idea':run('idea')}
    except Exception as e: return {'error':repr(e),'fallback':'neural comparison unavailable'}

def rollout(n=30000):
    ps=[]
    for s in np.linspace(-.2,.2,9):
        ok=0
        for _ in range(n//9):
            y=s
            for _ in range(300):
                if y<=U_EDGE: break
                if y>=T_EDGE: ok+=1; break
                y+=mu+sigma*rng.normal()
        ps.append(ok/(n//9))
    return ps

def main():
    h=exact_failure(); sweep=margin_sweep(); truth,mc,slope=mc_scaling(); nn=neural_comparison(h)
    d1,d2=sweep[1],sweep[2]
    slope_delta=(d2['min_initial_B']-d1['min_initial_B'])/(d2['delta']-d1['delta'])
    feas=[r for r in sweep if r['feasible']]; _,B=oracle(feas[-1]['delta'])
    result={'seed':SEED,'system':{'mu':mu,'sigma':sigma,'U':U_EDGE,'T':T_EDGE},'predictions':{'margin_slope_observed':slope_delta,'margin_prediction':'positive approximately linear until infeasible','mc_slope_observed':slope,'mc_prediction':-.5,'feasibility_cliff_between_delta': [.03,.04]},'margin_sweep':sweep,'mc_true_expectation':truth,'mc_scaling':mc,'exact_failure_max_X0':float(h[X0].max()),'rollout_success_by_start':rollout(),'certificate_check':{'delta':feas[-1]['delta'],'max_B_X0':float(B[X0].max()),'max_exact_failure_X0':float(h[X0].max()),'conservative_gap':float(B[X0].max()-h[X0].max())},'neural_comparison':nn}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()