#!/usr/bin/env python3
"""Reproduction of Figure 2 of the SGLRW paper (arXiv 2602.15925v1):
multimodal univariate target, exact gradient corrupted by symmetric alpha-stable
noise (alpha = 1.5) of increasing scale, lr = 1e-1, parallel chains.

Samplers: SGLD, Clipped-SGLD (R = sqrt(2 delta_t)), and SGLRW *with the
implementation clipping* of sqrt(delta_t/2)|grad U| to 1 that the paper
introduces at the end of Section 4 -- i.e. exactly the implementation whose
stability Figure 2 asserts.  The un-clipped SGLRW definition is included as a
control.

Stability statistics (the paper's figure reports none):
  * frac_diverged  : |theta| > 1e3 or non-finite at the final step
  * W1             : 1-Wasserstein distance between the final sample set and the
                     true target (exact quantile coupling, 4001 quantiles)
  * KS             : Kolmogorov-Smirnov distance to the true target CDF
  * sd, p99_abs    : spread statistics of the final sample set
"""
from __future__ import annotations
import argparse, json, math, time
import numpy as np

MIXES = {
    "A": ((0.4, -2.0, 1.0), (0.6, 2.0, 1.0)),
    "B": ((0.35, -2.5, 1.0), (0.65, 1.5, 1.2)),
    "C": ((0.5, -3.0, 0.8), (0.5, 2.0, 1.5)),
}


def mix_logpdf_grad(theta, mix):
    """-grad U = grad log p for a Gaussian mixture."""
    ws = np.array([m[0] for m in mix])[:, None]
    mus = np.array([m[1] for m in mix])[:, None]
    sds = np.array([m[2] for m in mix])[:, None]
    z = (theta[None, :] - mus) / sds
    logc = np.log(ws / (sds * np.sqrt(2 * np.pi))) - 0.5 * z ** 2
    mx = logc.max(0)
    w = np.exp(logc - mx)
    s = w.sum(0)
    resp = w / s
    grad_logp = -np.sum(resp * z / sds, axis=0)
    return -grad_logp                      # grad U


def mix_cdf(x, mix):
    from math import erf
    out = np.zeros_like(x)
    for w, mu, sd in mix:
        out += w * 0.5 * (1 + np.vectorize(erf)((x - mu) / (sd * np.sqrt(2))))
    return out


def mix_quantiles(probs, mix):
    lo, hi = -60.0, 60.0
    q = np.zeros_like(probs)
    a = np.full_like(probs, lo)
    b = np.full_like(probs, hi)
    for _ in range(80):
        m = 0.5 * (a + b)
        c = mix_cdf(m, mix)
        a = np.where(c < probs, m, a)
        b = np.where(c < probs, b, m)
    return 0.5 * (a + b)


def stable_rvs(rng, size, alpha=1.5):
    """Chambers-Mallows-Stuck symmetric alpha-stable, unit scale."""
    U = rng.uniform(-np.pi / 2, np.pi / 2, size)
    W = rng.exponential(1.0, size)
    return (np.sin(alpha * U) / np.cos(U) ** (1.0 / alpha)
            * (np.cos(U - alpha * U) / W) ** ((1.0 - alpha) / alpha))


def run(sampler, scale, seed, mix="A", delta0=0.1, iters=2000, chains=4000,
        schedule="const", alpha=1.5):
    m = MIXES[mix]
    rng = np.random.default_rng(7000 + seed * 101 + int(scale) * 7 + hash(sampler) % 97)
    theta = rng.standard_normal(chains) * 3.0
    clip_hits = 0
    total = 0
    for t in range(iters):
        dt = delta0 * (1.0 + t) ** (-0.55) if schedule == "decay" else delta0
        g = mix_logpdf_grad(theta, m)
        if scale > 0:
            g = g + scale * stable_rvs(rng, theta.shape, alpha)
        if sampler == "sgld":
            theta = theta - dt * g + np.sqrt(2 * dt) * rng.standard_normal(chains)
        elif sampler == "clipped_sgld":
            R = np.sqrt(2 * dt)
            theta = theta - np.clip(dt * g, -R, R) + R * rng.standard_normal(chains)
        elif sampler in ("sglrw", "sglrw_unclipped"):
            u = np.sqrt(dt / 2.0) * g
            clip_hits += int(np.sum(np.abs(u) > 1.0))
            total += chains
            if sampler == "sglrw":
                u = np.clip(u, -1.0, 1.0)
                p = 0.5 - 0.5 * u
            else:
                p = np.clip(0.5 - 0.5 * u, 0.0, 1.0)
            s = np.where(rng.random(chains) < p, 1.0, -1.0)
            theta = theta + np.sqrt(2 * dt) * s
        else:
            raise ValueError(sampler)
        theta = np.where(np.isfinite(theta), theta, np.sign(theta) * 1e12)
        theta = np.clip(theta, -1e12, 1e12)
    bad = ~np.isfinite(theta) | (np.abs(theta) > 1e3)
    good = theta[~bad]
    probs = (np.arange(4001) + 0.5) / 4001
    tq = mix_quantiles(probs, m)
    if good.size > 100:
        sq = np.quantile(good, probs)
        w1 = float(np.mean(np.abs(sq - tq)))
        grid = np.linspace(-15, 15, 3001)
        emp = np.searchsorted(np.sort(good), grid) / good.size
        ks = float(np.max(np.abs(emp - mix_cdf(grid, m))))
    else:
        w1, ks = float("inf"), 1.0
    return {"sampler": sampler, "scale": scale, "seed": seed, "mix": mix,
            "schedule": schedule, "delta0": delta0, "iters": iters, "chains": chains,
            "frac_diverged": float(bad.mean()),
            "W1": w1, "KS": ks,
            "sd": float(np.std(good)) if good.size > 100 else float("inf"),
            "p99_abs": float(np.quantile(np.abs(good), 0.99)) if good.size > 100 else float("inf"),
            "max_abs": float(np.max(np.abs(theta))),
            "clip_fraction": (clip_hits / total) if total else None}


def _job(j):
    return run(**j)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--quick", action="store_true")
    a = ap.parse_args()
    scales = [0.0, 1.0, 5.0, 10.0, 25.0, 50.0]
    samplers = ["sgld", "sglrw", "clipped_sgld"]
    seeds = range(2) if a.quick else range(5)
    jobs = []
    for sc in scales:
        for s in samplers:
            for sd in seeds:
                jobs.append(dict(sampler=s, scale=sc, seed=sd))
    for sc in [0.0, 25.0, 50.0]:
        for sd in seeds:
            jobs.append(dict(sampler="sglrw_unclipped", scale=sc, seed=sd))
        for s in samplers:
            for sd in seeds:
                jobs.append(dict(sampler=s, scale=sc, seed=sd, schedule="decay"))
            for mx in ["B", "C"]:
                for sd in seeds:
                    jobs.append(dict(sampler=s, scale=sc, seed=sd, mix=mx))
    from concurrent.futures import ProcessPoolExecutor
    out = []
    t0 = time.time()
    with ProcessPoolExecutor(max_workers=6) as ex:
        for i, r in enumerate(ex.map(_job, jobs)):
            out.append(r)
            if i % 20 == 0:
                print(i, len(jobs), round(time.time() - t0, 1), flush=True)
    with open(a.out, "w") as f:
        json.dump(out, f, indent=1)
    print("DONE", len(out), round(time.time() - t0, 1))


if __name__ == "__main__":
    main()