import json, random, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED=2130
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

def affine_action(p, bu=1., bd=.7, U=1., D=1.):
    return -U*np.sign(p*bu), D*np.sign(p*bd)

def rollout(grad, n=2000, steps=25, dt=.01, gamma=1., sigma=.02, seed=0):
    r=np.random.default_rng(seed); x=r.uniform(-1,1,n); occ=[]
    for _ in range(steps):
        x += -gamma*np.sign(grad(x))*dt + np.sqrt(2*sigma*sigma*dt)*r.normal(size=n)
        x=np.clip(x,-1,1); occ.append(x.copy())
    return np.concatenate(occ)

def ess_bins(x, bins=40):
    h=np.histogram(x,bins=bins,range=(-1,1))[0]; p=h[h>0]/h.sum()
    return float(1/(p*p).sum()), int((h>0).sum())

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self,z): return self.net(z)

def residual(m,z):
    z=z.requires_grad_(True); v=m(z); g=torch.autograd.grad(v.sum(),z,create_graph=True)[0]
    return g[:,1]+torch.abs(g[:,0])

def train(mode, seed):
    torch.manual_seed(seed); r=np.random.default_rng(seed); m=MLP(); opt=torch.optim.Adam(m.parameters(),lr=2e-3)
    uni=r.uniform([-1,0],[1,1],(800,2)).astype('float32'); traj=None
    for it in range(60):
        if mode=='trajectory' and (traj is None or it%15==0):
            xs=r.uniform(-1,1,100); ts=r.uniform(0,1,100)
            z=torch.tensor(np.c_[xs,ts],dtype=torch.float32,requires_grad=True)
            p=torch.autograd.grad(m(z).sum(),z)[0][:,0].detach().numpy()
            xs=np.clip(xs-.08*np.sign(p)+.03*r.normal(size=100),-1,1)
            traj=np.c_[np.repeat(xs,3),r.uniform(0,1,300)].astype('float32')
        if mode=='uniform': batch=uni[r.integers(0,len(uni),32)]
        else:
            a=traj[r.integers(0,len(traj),22)]; b=uni[r.integers(0,len(uni),10)]; batch=np.r_[a,b]
        z=torch.tensor(batch,dtype=torch.float32); loss=(residual(m,z)**2).mean()
        xb=torch.tensor(r.uniform(-1,1,8),dtype=torch.float32); target=torch.abs(xb)
        loss=loss+.5*((m(torch.stack([xb,torch.ones(8)],1)).squeeze()-target)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    test=torch.tensor(uni,dtype=torch.float32); rr=[]
    for q in test.split(200): rr.append(residual(m,q).detach().abs())
    rr=torch.cat(rr).numpy(); return float((rr**2).mean()),float(np.sort(rr)[-80:].dot(np.sort(rr)[-80:])/80)

def main():
    out={}
    # Prediction 1: affine min/max is attained at the sign-selected box corner.
    r=np.random.default_rng(1); errors=[]
    for p in r.normal(size=1000):
        u,d=affine_action(p); chosen=p*(1.3+u+.7*d)
        nested=min(p*(1.3+u0+.7*(D0)) for u0 in (-1,1) for D0 in ([-1,1] if p*.7 >= 0 else [1,-1]))
        # equivalent explicit min_u max_d, with d selected to maximize p*.7*d
        nested=min(max(p*(1.3+u0+.7*d0) for d0 in (-1,1)) for u0 in (-1,1))
        errors.append(abs(chosen-nested))
    out['affine_corner_check']={'max_abs_error':float(max(errors)),'prediction':'zero error'}
    # Prediction 2: deterministic steering displacement is gamma*dt.
    gs=[.25,.5,1.,2.]; dt=.01; measured=[]
    for g in gs:
        x=np.ones(10000); measured.append(float(np.mean(x-(x-g*dt))))
    out['gamma_scaling']={'gamma':gs,'observed':measured,'predicted':[g*dt for g in gs], 'ratios':[measured[i]/(gs[i]*dt) for i in range(4)]}
    # Prediction 3: Brownian increment variance is 2*sigma^2*dt.
    ss=[.02,.05,.1]; rv=np.random.default_rng(2); observed=[]
    for s in ss: observed.append(float(np.var(np.sqrt(2*s*s*dt)*rv.normal(size=200000))))
    out['sigma_scaling']={'sigma':ss,'observed_variance':observed,'predicted':[2*s*s*dt for s in ss], 'ratios':[observed[i]/(2*ss[i]**2*dt) for i in range(3)]}
    # Prediction 4: mixture has exactly alpha trajectory samples and retains reservoir samples.
    tr=rollout(lambda x:x,n=1000,steps=10,dt=.01,gamma=1,sigma=.02,seed=3); u=rv.uniform(-1,1,len(tr)); mix=[]
    for a in (0,.25,.5,.75,1.):
        k=int(a*len(tr)); x=np.r_[tr[:k],u[:len(tr)-k]]; mix.append([a,len(x),k/len(x),*ess_bins(x)])
    out['mixture_check']={'columns':['alpha','N','observed_traj_fraction','ESS','occupied_bins'],'rows':mix}
    t=time.time(); base=[train('uniform',s) for s in (0,1)]; idea=[train('trajectory',s) for s in (0,1)]
    out['pinn']={'uniform_mean':np.mean(base,axis=0).tolist(),'trajectory_mean':np.mean(idea,axis=0).tolist(),'uniform_runs':base,'trajectory_runs':idea,'seconds':time.time()-t}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()