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

SEED = 2675
rng = np.random.default_rng(SEED)
OUT = Path('results.csv')

# Harmonic Langevin chain: x'=(1-eta*lambda_k)x + N(0,2*T*eta).
def log_normal(y, mean, var):
    return -0.5 * (math.log(2*math.pi*var) + (y-mean)**2/var)

def path_ep(x, lam, eta=0.08, T=1.0):
    # p0 and pK^R are equilibrium endpoint densities.
    lp0 = log_normal(x[0], 0., T/lam[0])
    lpk = log_normal(x[-1], 0., T/lam[-1])
    lf, lr = lp0, lpk
    for k in range(len(lam)-1):
        a=(1-eta*lam[k])
        vf = (T/lam[k])*(1-a*a)
        mf = a*x[k]
        # reverse transition at reverse time k uses u^R_k = u_{K-1-k}
        # evaluated on (x[k+1] -> x[k]); this is the paired forward protocol value.
        rr=len(lam)-2-k
        ar=(1-eta*lam[rr])
        mr = ar*x[k+1]
        lf += log_normal(x[k+1], mf, vf)
        lr += log_normal(x[k], mr, vf)
    return lf-lr

def simulate(lam, n=30000, eta=0.08, T=1.0):
    K=len(lam)-1
    x=np.empty((n,K+1))
    x[:,0]=rng.normal(0, math.sqrt(T/lam[0]), size=n)
    for k in range(K):
        a=1-eta*lam[k]
        var=(T/lam[k])*(1-a*a)
        x[:,k+1]=a*x[:,k]+rng.normal(0, math.sqrt(var), size=n)
    lf=log_normal(x[:,0], 0., T/lam[0])
    lr=log_normal(x[:,-1], 0., T/lam[-1])
    for k in range(K):
        af=1-eta*lam[k]
        vf=(T/lam[k])*(1-af*af)
        mf=af*x[:,k]
        ar=1-eta*lam[K-1-k]
        vf_r=(T/lam[K-1-k])*(1-ar*ar)
        mr=ar*x[:,k+1]
        lf += log_normal(x[:,k+1],mf,vf)
        lr += log_normal(x[:,k],mr,vf_r)
    vals=lf-lr
    return float(np.mean(vals)), float(np.std(vals)/math.sqrt(n))

def scaling_sweeps():
    # Prediction 1: equilibrium constant protocol has zero mean EP (within MC error).
    constant=np.ones(41)*1.5
    zero, zse=simulate(constant, n=12000)
    # Prediction 2: near-equilibrium protocol EP is quadratic in amplitude.
    amps=np.array([.02,.04,.08,.16,.32])
    amp_ep=[]
    for a in amps:
        lam=1.5 + a*np.linspace(-1,1,41)
        amp_ep.append(max(simulate(lam,n=16000)[0],1e-10))
    amp_slope=float(np.polyfit(np.log(amps),np.log(amp_ep),1)[0])
    # Prediction 3: slow-ramp irreversible EP decreases approximately as ramp speed^2
    # when duration is held in physical units via eta and K is varied.
    Ks=np.array([10,20,40,80,160])
    rate_ep=[]
    for K in Ks:
        lam=1.5 + .25*np.linspace(-1,1,K+1)
        rate_ep.append(max(simulate(lam,n=14000)[0],1e-10))
    # eps is per-step protocol change.
    eps=.5/Ks
    rate_slope=float(np.polyfit(np.log(eps),np.log(rate_ep),1)[0])
    return {'constant_mean_ep':zero,'constant_mc_se':zse,'amplitudes':amps.tolist(),
            'amplitude_ep':amp_ep,'amplitude_loglog_slope':amp_slope,
            'ramp_K':Ks.tolist(),'ramp_eps':eps.tolist(),'ramp_ep':rate_ep,
            'ramp_loglog_slope':rate_slope}

# Tiny quadratic regression, with exact optimizer transition likelihood for stored states.
def run_optimizer(adaptive, seed=2675, steps=500):
    r=np.random.default_rng(seed)
    n,d=256,4
    X=r.normal(size=(n,d)); true=r.normal(size=d); y=X@true+r.normal(scale=.15,size=n)
    w=r.normal(scale=.2,size=d); T=.015; eta=.035; target=.08; alpha=.10
    losses=[]; sigmas=[]; window=30; states=[]; grads=[]
    for t in range(steps):
        # deterministic full-batch gradient keeps the transition model explicit.
        g=X.T@(X@w-y)/n
        old=w.copy(); noise=r.normal(size=d)
        w=w-eta*g+math.sqrt(2*T*eta)*noise
        states.append(old); grads.append(g)
        if len(states)>=window:
            # use current window's actual eta schedule as a constant protocol;
            # reverse likelihood uses gradients recomputed at endpoint states.
            xx=np.array(states[-window:]+[w.copy()]); gg=np.array(grads[-window:])
            # endpoint-to-endpoint Gaussian path ratio (p endpoints cancel approximately).
            sig=0.
            for j in range(window):
                mf=xx[j]-eta*gg[j]
                vf=2*T*eta
                gr=X.T@(X@xx[j+1]-y)/n
                mr=xx[j+1]-eta*gr
                sig += np.sum((xx[j+1]-mf)**2-(xx[j]-mr)**2)/(2*vf)
            sig/=window
            sigmas.append(sig)
            if adaptive:
                eta=float(np.clip(eta*math.exp(-alpha*(sig-target)),.008,.07))
        losses.append(float(np.mean((X@w-y)**2)))
    return {'final_loss':float(np.mean(losses[-50:])),'best_loss':float(np.min(losses)),
            'mean_last_entropy':float(np.mean(sigmas[-100:])) if sigmas else 0.,
            'eta_final':eta,'loss_trace':losses}

def main():
    sweep=scaling_sweeps()
    base=run_optimizer(False); idea=run_optimizer(True)
    rows=[]
    for a,e in zip(sweep['amplitudes'],sweep['amplitude_ep']): rows.append(['amplitude',a,e])
    for k,e in zip(sweep['ramp_K'],sweep['ramp_ep']): rows.append(['ramp_K',k,e])
    with OUT.open('w',newline='') as f:
        w=csv.writer(f); w.writerow(['sweep','parameter','mean_entropy_production']); w.writerows(rows)
    out={'seed':SEED,'toy':sweep,'baseline':{k:v for k,v in base.items() if k!='loss_trace'},
         'entropy_controller':{k:v for k,v in idea.items() if k!='loss_trace'}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()