import json, math, time
import numpy as np

SEED = 2485
D, K = 20, 3
rng0 = np.random.default_rng(SEED)
one = rng0.normal(0, 0.7, size=(D, K))
pair = rng0.normal(0, 0.06, size=(D-1, K, K))

def energy(x):
    v = float(np.sum(one[np.arange(D), x]))
    for i in range(D-1): v += pair[i, x[i], x[i+1]]
    return v

def rates(x, p):
    u0 = energy(x); q = np.zeros((D, K))
    for i in range(D):
        for z in range(K):
            if z != x[i]:
                y=x.copy(); y[i]=z
                q[i,z]=max(float(p[i]),0.0)*math.exp(-0.5*(energy(y)-u0))
    return q

def exact_path(x, p, rng, horizon):
    t=0.; n=0
    while t < horizon:
        q=rates(x,p); lam=float(q.sum())
        if lam <= 0: break
        t += rng.exponential(1./lam)
        if t > horizon: break
        flat=int(rng.choice(D*K,p=(q.ravel()/lam)))
        i,z=divmod(flat,K); x[i]=z; n+=1
    return n

def tau_path(x, p, rng, horizon, h):
    n=0; t=0.
    while t < horizon-1e-12:
        step=min(h,horizon-t); q=rates(x,p)
        active=rng.random(q.shape) < (-np.expm1(-step*q))
        for i in range(D):
            zs=np.flatnonzero(active[i])
            if len(zs):
                w=q[i,zs]; x[i]=int(zs[rng.choice(len(zs),p=w/w.sum())]); n+=1
        t += step
    return n

def make_state(rng): return rng.integers(0,K,size=D), rng.normal(0,1,size=D)

def poisson_checks():
    q=0.73; hvals=[0.05,0.2,0.8,1.5]; N=300000; rows=[]
    for h in hvals:
        s=np.random.default_rng(SEED+int(100*h)).poisson(h*q,N)
        rows.append({'h':h,'mean_obs':float(s.mean()),'mean_pred':h*q,
                     'nonzero_obs':float(np.mean(s>0)),'nonzero_pred':float(1-math.exp(-h*q))})
    return rows

def bias_sweep():
    # Compare endpoint distributions to exact trajectories using common initial states.
    hs=[0.01,0.025,0.05,0.1,0.2]
    M=120; horizon=.25; out=[]
    for h in hs:
        exact=np.zeros((M,D),int); approx=np.zeros((M,D),int); ne=na=0
        for m in range(M):
            a,b=make_state(np.random.default_rng(SEED+m));
            xe=a.copy(); xa=a.copy();
            ne+=exact_path(xe,b,np.random.default_rng(10000+m),horizon)
            na+=tau_path(xa,b,np.random.default_rng(20000+m),horizon,h)
            exact[m]=xe; approx[m]=xa
        # Total variation of coordinate-0 marginals and mismatch rate.
        tv=0.
        for z in range(K):
            tv += abs(np.mean(exact[:,0]==z)-np.mean(approx[:,0]==z))
        tv*=.5
        out.append({'h':h,'coord0_TV':float(tv),'endpoint_mismatch':float(np.mean(np.any(exact!=approx,axis=1))),
                    'exact_events_per_path':ne/M,'tau_events_per_path':na/M})
    return out

def throughput():
    M=100; horizon=.25; h=.05
    states=[]; moms=[]
    for m in range(M):
        x,p=make_state(np.random.default_rng(50000+m)); states.append(x); moms.append(p)
    t=time.perf_counter(); en=0
    for m in range(M): en+=exact_path(states[m].copy(),moms[m],np.random.default_rng(60000+m),horizon)
    te=time.perf_counter()-t
    t=time.perf_counter(); tn=0
    for m in range(M): tn+=tau_path(states[m].copy(),moms[m],np.random.default_rng(70000+m),horizon,h)
    tt=time.perf_counter()-t
    return {'M':M,'h':h,'exact_sec':te,'tau_sec':tt,'exact_paths_per_sec':M/te,'tau_paths_per_sec':M/tt,'speedup':te/tt,'exact_events':en/M,'tau_events':tn/M}

def main():
    result={'seed':SEED,'dimensions':[D,K],'poisson_checks':poisson_checks(),
            'bias_sweep':bias_sweep(),'throughput':throughput()}
    # Fit log-log observed bias scaling, excluding numerical floor if necessary.
    hs=np.array([r['h'] for r in result['bias_sweep']]); tv=np.array([r['coord0_TV'] for r in result['bias_sweep']])
    result['loglog_bias_slope']=float(np.polyfit(np.log(hs),np.log(np.maximum(tv,1e-8)),1)[0])
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()