import json, math, time
import numpy as np

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


def clip(v, bound):
    n = np.linalg.norm(v)
    return v if n <= bound else v * (bound / n)


def ar1(n, rho, rng, burn=300):
    x = np.zeros(n + burn)
    noise = rng.normal(size=n + burn)
    for i in range(1, len(x)):
        x[i] = rho * x[i-1] + math.sqrt(1-rho*rho) * noise[i]
    return x[burn:]


def means(z, b):
    return np.array([z[i:i+b].mean() for i in range(0, len(z)-b+1, b)])


def correction_samples(rho, level, trials=6000, b0=4):
    b = b0 * (2**level)
    # Vectorized independent trajectories; each row is one coupled block.
    e = rng.normal(size=(trials, b))
    z = np.zeros_like(e)
    z[:, 0] = e[:, 0]
    scale = math.sqrt(1-rho*rho)
    for i in range(1, b):
        z[:, i] = rho*z[:, i-1] + scale*e[:, i]
    vals = z.mean(axis=1) - z[:, :b//2].mean(axis=1)
    return vals


def telescoping_check():
    # For p_l proportional to 2^-l, E[g0 + Delta_l/p_l] = E[g_K].
    K, b0, rho, trials = 4, 4, .8, 30000
    p = np.array([2.0**(-l) for l in range(K+1)])
    p /= p.sum()
    est, fine = [], []
    for _ in range(trials):
        z = ar1(b0*2**K, rho, rng)
        gs = [z[:b0*2**l].mean() for l in range(K+1)]
        l = rng.choice(K+1, p=p)
        est.append(gs[0] if l == 0 else gs[0] + (gs[l]-gs[l-1])/p[l])
        fine.append(gs[K])
    est, fine = np.array(est), np.array(fine)
    diff = est.mean() - fine.mean()
    se = math.sqrt(est.var()/trials + fine.var()/trials)
    return {"observed_difference": float(diff), "two_se": float(2*se),
            "passes_2se": bool(abs(diff) <= 2*se)}


def math_sweeps():
    # Prediction A: correction variance is proportional to AR long-run variance
    # (1+rho)/(1-rho), and shrinks with level approximately as 1/b.
    rows = []
    for rho in [0.0, .5, .8, .95]:
        vals = correction_samples(rho, 3)
        rows.append({"rho": rho, "var_delta_level3": float(vals.var()),
                     "predicted_dependence_multiplier": (1+rho)/(1-rho)})
    # Compare observed ratios against the iid baseline.
    base = rows[0]["var_delta_level3"]
    for r in rows: r["observed_ratio_to_iid"] = r["var_delta_level3"] / base
    levels = []
    for level in range(5):
        v = correction_samples(.8, level, b0=32).var()
        levels.append({"level": level, "block": 32*2**level, "variance": float(v),
                       "variance_times_block": float(v*(4*2**level))})
    # Fixed-large-block prediction: Var(mean) ratio tends to (1+rho)/(1-rho).
    dep_fixed = []
    bfix, trials_fix = 512, 4000
    for rho in [0.0, .5, .8, .95]:
        e = rng.normal(size=(trials_fix, bfix)); z = np.zeros_like(e); z[:, 0] = e[:, 0]
        q = math.sqrt(1-rho*rho)
        for i in range(1, bfix): z[:, i] = rho*z[:, i-1] + q*e[:, i]
        dep_fixed.append({"rho": rho, "block": bfix,
                          "observed_mean_variance_ratio": float(z.mean(1).var() / dep_fixed[0]["_var"] if dep_fixed else 1.0),
                          "predicted_ratio": (1+rho)/(1-rho), "_var": float(z.mean(1).var())})
    for r in dep_fixed: r.pop("_var", None)
    # Prediction B: clipping enforces the pathwise norm bound exactly.
    max_norm = 0.0; clipped = 0
    for _ in range(10000):
        v = rng.normal(size=7)*10
        c = clip(v, 1.0)
        max_norm = max(max_norm, np.linalg.norm(c))
        clipped += np.linalg.norm(v) > 1
    return {"dependence": rows, "fixed_large_block_dependence": dep_fixed, "level_scaling_rho_0.8": levels,
            "clipping": {"bound": 1.0, "max_observed_norm": float(max_norm),
                          "clipping_fraction": clipped/10000}}


def grad(theta, z):
    # Two-parameter linear neural model, squared loss gradient, ordered stream.
    x, y = z
    e = theta[0]*x + theta[1] - y
    return np.array([e*x, e])


def stream_data(n, rho, rng):
    x = ar1(n, rho, rng)
    # stationary regression target with independent observation noise
    y = 1.7*x - .35 + .1*rng.normal(size=n)
    return np.stack([x, y], axis=1)


def train(method, rho, seed=1414, steps=350, b0=4, K=3, lr=.08):
    rg = np.random.default_rng(seed)
    # Equal expected sample budget: baseline consumes b0*2^K, MLMC expected cost
    # is b0 * sum p_l 2^l, and we use a fixed reserved fine block for simplicity.
    data = stream_data(steps*b0*2**K + 20, 0.0 if method == "iid" else rho, rg)
    th = np.array([0., 0.]); losses=[]; cursor=0; gradseq=[]; clip_count=0
    p = np.array([2.**(-l) for l in range(K+1)]); p /= p.sum()
    for _ in range(steps):
        if method == 'baseline':
            block = data[cursor:cursor+b0*2**K]; cursor += b0*2**K
            gh = np.mean([grad(th,z) for z in block], axis=0)
        elif method == 'coupled':
            block = data[cursor:cursor+b0*2**K]; cursor += b0*2**K
            gs=[]
            for l in range(K+1):
                b=b0*2**l; gs.append(np.mean([grad(th,z) for z in block[:b]],axis=0))
            l=rg.choice(K+1,p=p)
            gh=gs[0] if l==0 else gs[0]+(gs[l]-gs[l-1])/p[l]
        elif method == 'iid':
            block = data[cursor:cursor+b0*2**K].copy(); cursor += b0*2**K
            rg.shuffle(block)
            gh=np.mean([grad(th,z) for z in block],axis=0)
        n=np.linalg.norm(gh)
        if n>4: gh=gh*4/n; clip_count += 1
        th -= lr*gh
        gradseq.append(gh)
        # validation risk on population proxy
        losses.append((th[0]-1.7)**2 + (th[1]+.35)**2)
    a=np.array(gradseq); ac=np.corrcoef(a[:-1,0],a[1:,0])[0,1]
    return {"final_parameter_error":float(losses[-1]), "loss_at_100":float(losses[99]),
            "loss_at_350":float(losses[-1]), "gradient_lag1_autocorr":float(ac),
            "clip_fraction":clip_count/steps, "samples":int(cursor)}


def main():
    out={"seed":SEED, "telescoping":telescoping_check(), "math_sweeps":math_sweeps(),
         "benchmark":{}}
    for rho in [0.0,.8,.95]:
        out["benchmark"][str(rho)]={m:train(m,rho) for m in ["baseline","coupled","iid"]}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()