import json, math
import numpy as np
from scipy.special import ndtr

SEED = 1729
rng = np.random.default_rng(SEED)

def f(x):
    return 0.85*x + 0.20*np.tanh(x)

def jac(x):
    return 0.85 + 0.20/(np.cosh(x)**2)

class Comp:
    __slots__ = ('w','m','v')
    def __init__(self, w, m, v): self.w, self.m, self.v = float(w), float(m), float(v)

def normalize(cs):
    s = sum(c.w for c in cs)
    for c in cs: c.w /= s
    return cs

def gaussian_w2(a, b):
    return math.sqrt((a.m-b.m)**2 + (math.sqrt(max(a.v,0))-math.sqrt(max(b.v,0)))**2)

def merge(cs, eps):
    cs = list(cs)
    while len(cs) > 1:
        cand = [(gaussian_w2(cs[i],cs[j]),i,j) for i in range(len(cs)) for j in range(i+1,len(cs))]
        d,i,j = min(cand)
        if d > eps: break
        a,b = cs[i],cs[j]; w = a.w+b.w; m = (a.w*a.m+b.w*b.m)/w
        v = (a.w*(a.v+(a.m-m)**2)+b.w*(b.v+(b.m-m)**2))/w
        cs = [c for k,c in enumerate(cs) if k not in (i,j)] + [Comp(w,m,v)]
    return normalize(cs)

def propagate(cs, d, q, eps=None):
    out=[]
    for c in cs:
        J=jac(c.m)
        out += [Comp(.5*c.w, f(c.m)-d, J*J*c.v+q), Comp(.5*c.w, f(c.m)+d, J*J*c.v+q)]
    return merge(out,eps) if eps is not None else normalize(out)

def moments(cs):
    m=sum(c.w*c.m for c in cs)
    v=sum(c.w*(c.v+(c.m-m)**2) for c in cs)
    return m,v

def chance(cs, threshold):
    return sum(c.w*ndtr((threshold-c.m)/math.sqrt(c.v)) for c in cs)

def sample(cs, n, rg):
    ix=rg.choice(len(cs), n, p=[c.w for c in cs])
    return np.array([rg.normal(cs[k].m, math.sqrt(cs[k].v)) for k in ix])

def empirical_w2(x,y):
    # In one dimension the sorted empirical coupling is optimal W2.
    n=min(len(x),len(y)); xx=np.sort(x[:n]); yy=np.sort(y[:n])
    return math.sqrt(np.mean((xx-yy)**2))

def gaussian_component(cs):
    m,v=moments(cs)
    return [Comp(1,m,v)]

def main():
    q=.0025; d=.9; initial=[Comp(1,0,.01)]
    # Prediction 1: after one step, mode separation is 2d, independent of covariance.
    sep=[]
    for dd in [.1,.3,.6,.9,1.2]:
        z=propagate(initial,dd,q)
        sep.append({'d':dd,'observed':z[1].m-z[0].m,'predicted':2*dd})
    # Prediction 2: two equal Gaussian modes merge precisely at W2 tolerance >= separation.
    threshold_sweep=[]
    raw=propagate(initial,d,q)
    rawdist=gaussian_w2(raw[0],raw[1])
    for eps in [.2,.9,1.5,1.79,1.8,1.81,2.0]:
        threshold_sweep.append({'eps':eps,'components':len(merge(raw,eps)),'pair_w2':rawdist,'predicted_merge':eps >= rawdist})
    # Prediction 3: merging at tolerance above pair W2 changes the represented law,
    # while tolerances below it preserve both modes.
    merge_error = []
    rg = np.random.default_rng(SEED + 77)
    reference = sample(raw, 30000, rg)
    for eps in [.9, 1.79, 1.81, 2.0]:
        mc = merge(raw, eps)
        approx = sample(mc, 30000, np.random.default_rng(SEED + int(100*eps)))
        merge_error.append({'eps':eps,'components':len(mc),'sample_w2_to_unmerged':empirical_w2(reference,approx),
                            'predicted_merged':eps >= rawdist})
    # Prediction 4: before merging, chance probability is mixture CDF and differs from moment Gaussian.
    # A symmetric threshold at zero is exactly 0.5; use a positive threshold for a quantitative sweep.
    chance_rows=[]
    for dd in [.1,.3,.6,.9,1.2]:
        z=propagate(initial,dd,q); mm=gaussian_component(z)
        p_mix=chance(z,.5); p_gauss=chance(mm,.5)
        chance_rows.append({'d':dd,'mixture':p_mix,'moment_gaussian':p_gauss,'absolute_gap':abs(p_mix-p_gauss)})
    # Small nonlinear rollout comparison against moment matching, with exact mixture as reference.
    horizons=[1,2,4,6,8]
    rollout=[]
    for H in horizons:
        exact=initial
        mix=initial
        for _ in range(H):
            exact=propagate(exact,d,q)
            mix=propagate(mix,d,q,eps=1e9) # one component after every step
        rg=np.random.default_rng(SEED+H)
        a=sample(exact,20000,rg); b=sample(mix,20000,rg)
        rollout.append({'horizon':H,'components_exact':len(exact),'components_gaussian':len(mix),
                        'sample_w2_error_gaussian_vs_mixture':empirical_w2(a,b),
                        'exact_mean':moments(exact)[0],'gaussian_mean':moments(mix)[0]})
    out={'seed':SEED,'parameters':{'d':d,'q':q},'prediction_mode_separation':sep,
         'prediction_merge_transition':threshold_sweep,'prediction_chance_calibration':chance_rows,'merge_error_sweep':merge_error,
         'rollout_comparison':rollout}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()