import json, math
import numpy as np

SEED = 2817

def telegraph_run(gamma, lambdas, vs, dt=0.01, n=60000, burn=6000, seed=0):
    rng = np.random.default_rng(seed)
    lambdas, vs = np.asarray(lambdas, float), np.asarray(vs, float)
    s = rng.choice([-1., 1.], size=len(vs))
    u = 0.0
    a = math.exp(-gamma*dt)
    b = (1-a)/gamma
    p = -np.expm1(-lambdas*dt)
    out = np.empty(n-burn)
    maxabs = 0.
    j = 0
    for i in range(n):
        s[rng.random(len(s)) < p] *= -1
        u = a*u + b*float(np.dot(vs, s))
        maxabs = max(maxabs, abs(u))
        if i >= burn:
            out[j] = u; j += 1
    return out, maxabs

def variance_pred(gamma, lambdas, vs):
    return float(np.sum(np.asarray(vs)**2/(gamma*(gamma+2*np.asarray(lambdas)))))

def cov_correct(gamma, lam, v, tau):
    # Correct covariance of du=-gamma*u+v*s, E[s(t)s(t+tau)]=exp(-2 lam tau).
    tau = np.asarray(tau, float)
    if abs(gamma-2*lam) < 1e-8:
        # continuous limiting form
        return v*v*np.exp(-gamma*tau)*(1/(2*gamma) + tau/2)
    return v*v*(gamma*np.exp(-2*lam*tau)-2*lam*np.exp(-gamma*tau))/(gamma*(gamma**2-4*lam**2))

def cov_as_given(gamma, lam, v, tau):
    tau=np.asarray(tau,float)
    return v*v*(np.exp(-2*lam*tau)-np.exp(-gamma*tau))/(gamma**2-4*lam**2)

def acf(x, maxlag):
    x=x-x.mean(); den=np.dot(x,x)
    return np.array([np.dot(x[:-k] if k else x, x[k:] if k else x)/den for k in range(maxlag+1)])

def excess_kurtosis(x):
    z=x-x.mean(); m2=np.mean(z*z)
    return float(np.mean(z**4)/m2**2-3)

def quadratic_opt(noise_kind, dim=20, steps=5000, lr=.08, alpha=1., gamma=1., lam=.7, K=4, seed=1):
    rng=np.random.default_rng(seed); theta=rng.normal(2., .5, dim)
    u=np.zeros(dim); s=rng.choice([-1.,1.], size=(dim, K))
    # Match stationary per-coordinate telegraph variance with OU noise.
    v=1./math.sqrt(K)
    var=K*v*v/(gamma*(gamma+2*lam))
    ou=np.zeros(dim); oo=math.exp(-gamma)
    losses=[]
    for t in range(steps):
        g=theta
        if noise_kind=='telegraph':
            s[rng.random((dim,K)) < -np.expm1(-lam)] *= -1
            u=oo*u+(1-oo)*(v*np.sum(s,axis=1)/gamma)
            force=u
        elif noise_kind=='ou':
            ou=oo*ou+math.sqrt(var*(1-oo*oo))*rng.normal(size=dim)
            force=ou
        else: force=0.
        theta -= lr*(g+alpha*force)
        losses.append(float(np.mean(theta*theta)))
    return float(np.mean(losses[-500:])), float(np.mean(theta*theta))

def main():
    gamma=1.; dt=.01
    rows=[]
    # Prediction 1: hard bound independent of lambda and time.
    for lam in [.1, .7, 2., 8.]:
        x,m=telegraph_run(gamma,[lam]*4,[.25]*4,dt=dt,seed=int(lam*100))
        rows.append({'lambda':lam,'bound':1.0,'max_abs':m,'ratio':m})
    # Prediction 2: exact variance scaling in lambda.
    variance=[]
    for lam in [.1,.3,.7,1.5,3.,8.]:
        x,_=telegraph_run(gamma,[lam]*4,[.25]*4,dt=dt,seed=10+int(lam*10))
        pred=variance_pred(gamma,[lam]*4,[.25]*4)
        variance.append({'lambda':lam,'emp_var':float(np.var(x)),'pred_var':pred,'rel_err':float(np.var(x)/pred-1)})
    # Prediction 3: covariance is sum of exponentials; compare corrected derivation and supplied expression.
    lam=.7; x,_=telegraph_run(gamma,[lam],[1.],dt=dt,seed=77)
    lags=np.array([0,10,25,50,100,200,400]); empirical=acf(x,int(lags[-1]))[lags]
    empirical_cov=empirical*np.var(x)
    corrected=cov_correct(gamma,lam,1.,lags*dt)
    supplied=cov_as_given(gamma,lam,1.,lags*dt)
    cov={'lags':(lags*dt).tolist(),'empirical':empirical_cov.tolist(),'corrected':corrected.tolist(),'supplied':supplied.tolist(),
         'rmse_corrected':float(np.sqrt(np.mean((empirical_cov-corrected)**2))),
         'rmse_supplied':float(np.sqrt(np.mean((empirical_cov-supplied)**2)))}
    # Prediction 4: many independent sources reduce standardized excess kurtosis.
    kurt=[]
    for K in [1,2,4,8,16]:
        x,_=telegraph_run(gamma,[.7]*K,[1/math.sqrt(K)]*K,dt=dt,seed=100+K)
        kurt.append({'K':K,'excess_kurtosis':excess_kurtosis(x),'variance':float(np.var(x)),
                     'pred_var':variance_pred(gamma,[.7]*K,[1/math.sqrt(K)]*K)})
    # Small secondary optimization sanity check (one telegraph source per coordinate).
    opt={k:quadratic_opt(k,seed=200) for k in ['none','telegraph','ou']}
    result={'config':{'gamma':gamma,'dt':dt,'steps':60000,'burn':30000},'bound_sweep':rows,
            'variance_sweep':variance,'covariance_check':cov,'kurtosis_sweep':kurt,'quadratic_optimizer':opt,
            'note':'The covariance formula in the prompt is tested literally; it is zero at tau=0 and is not a valid covariance. corrected is derived from the same SDE.'}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()