import json, math, os
import numpy as np

# Inverse-square reset diffusion toy experiment.
# Units: dt is the diffusion time step and D=1 unless specified.
def beta_pred(ratio):
    return (math.sqrt(1.0 + 4.0 * ratio) - 1.0) / 2.0

def simulate(r0, lam=-2.0, D=1.0, eps=0.1, dt=0.002, nwalk=1800,
             steps=9000, burn=1500, seed=0, box=80.0):
    rng = np.random.default_rng(seed)
    z = np.zeros(nwalk, dtype=np.float64)
    samples = []
    reset_count = 0
    for t in range(steps):
        z += math.sqrt(2*D*dt) * rng.standard_normal(nwalk)
        # generalized hazard; for lambda=-2 this is exactly the proposal
        if lam == -2:
            rate = r0 / (z*z + eps*eps)
        else:
            rate = r0 / (np.abs(z) + eps) ** (-lam)
        q = 1.0 - np.exp(-dt * rate)
        hit = rng.random(nwalk) < q
        reset_count += int(hit.sum())
        z[hit] = 0.0
        # A large reflecting box avoids numerical overflow; box dependence is measured separately.
        z[z > box] = 2*box - z[z > box]
        z[z < -box] = -2*box - z[z < -box]
        if t >= burn and t % 10 == 0:
            samples.append(z.copy())
    x = np.abs(np.concatenate(samples))
    return x, reset_count / (nwalk * steps)

def tail_fit(x, lo=2.0, hi=15.0):
    # CCDF ~ x^-(beta-1), so density exponent is one minus CCDF slope.
    xx = np.linspace(lo, hi, 40)
    cc = np.array([(x >= a).mean() for a in xx])
    good = (cc > 3e-4) & (cc < .3)
    if good.sum() < 5: return float('nan'), float('nan')
    slope, intercept = np.polyfit(np.log(xx[good]), np.log(cc[good]), 1)
    return 1.0 - slope, float(np.corrcoef(np.log(xx[good]), np.log(cc[good]))[0,1])

def run_toy():
    ratios = [0.5, 1.0, 1.5, 2.0, 2.5, 4.0, 8.0]
    rows=[]
    for i,r in enumerate(ratios):
        x, rr = simulate(r, seed=100+i)
        b, fitq = tail_fit(x)
        rows.append({'r0_over_D':r, 'beta_pred':beta_pred(r), 'beta_fit':b,
                     'tail_fit_r':fitq, 'mean_abs':float(x.mean()),
                     'p_abs_gt_5':float((x>5).mean()), 'reset_rate_step':rr})
    # Boundary prediction: above ratio 2 has finite first moment only for beta>2,
    # while normalizability itself is beta>1. Compare probability mass in a growing box.
    box_rows=[]
    for r in [1.0, 4.0]:
        vals=[]
        for B in [10.,20.,40.]:
            x,_=simulate(r, box=B, steps=8000, burn=1200, seed=900+int(B)+int(r*10))
            vals.append({'box':B,'mean_abs':float(x.mean()),'p_gt_5':float((x>5).mean())})
        box_rows.append({'ratio':r,'beta_pred':beta_pred(r),'box_scaling':vals})
    # Controls for the claimed lambda classification, at r0=4.
    controls=[]
    for lam in [-3.,-2.,-1.]:
        x,_=simulate(4.,lam=lam,steps=9000,burn=1500,seed=700+int((lam+3)*10))
        b,q=tail_fit(x)
        controls.append({'lambda':lam,'mean_abs':float(x.mean()),'tail_density_exponent_fit':b,'fit_r':q})
    return {'toy_rows':rows,'box_rows':box_rows,'lambda_controls':controls}

def run_ml(seed=42):
    # Small, reproducible regression; compare SGD with a checkpoint reset wrapper.
    import torch
    torch.manual_seed(seed); np.random.seed(seed)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        dev=torch.device(device)
        x=torch.linspace(-3,3,512,device=dev).unsqueeze(1)
        y=torch.sin(3*x)+0.15*torch.randn_like(x)
        tr=x[:384]; ty=y[:384]; va=x[384:]; vy=y[384:]
        def train(reset=False):
            torch.manual_seed(seed)
            m=torch.nn.Sequential(torch.nn.Linear(1,32),torch.nn.Tanh(),torch.nn.Linear(32,1)).to(dev)
            opt=torch.optim.SGD(m.parameters(),lr=.035,momentum=.8)
            ref=[p.detach().clone() for p in m.parameters()]
            best=1e9; resets=0
            for step in range(700):
                idx=torch.randint(0,len(tr),(64,),device=dev)
                opt.zero_grad(); loss=((m(tr[idx])-ty[idx])**2).mean(); loss.backward(); opt.step()
                with torch.no_grad():
                    vl=((m(va)-vy)**2).mean().item()
                    if vl<best: best=vl; ref=[p.detach().clone() for p in m.parameters()]
                    if reset:
                        z=math.sqrt(sum(((p-r)**2).sum().item() for p,r in zip(m.parameters(),ref)))
                        # D estimate is deliberately fixed to the toy scale; this is a stress test, not tuning.
                        q=1-math.exp(-.004/(z*z+.05*.05))
                        if np.random.random()<q:
                            for p,r in zip(m.parameters(),ref): p.copy_(r)
                            opt.state.clear(); resets+=1
            return float(((m(va)-vy)**2).mean().item()),resets
        base=train(False); idea=train(True)
        return {'device':str(dev),'baseline_val_mse':base[0],'reset_val_mse':idea[0],'reset_count':idea[1]}
    except Exception as e:
        return {'device':'cpu-fallback','error':repr(e)}

if __name__=='__main__':
    out={'toy':run_toy(),'ml':run_ml()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))