import json, math
from pathlib import Path
import numpy as np
from scipy.special import logsumexp

SEED = 3134
rng = np.random.default_rng(SEED)
D, MODES, M = 16, 8, 1200
SIGMA = 0.55
LOG2PI = math.log(2.0 * math.pi)

# Eight separated modes, with a common translation changing over annealing levels.
base_means = np.zeros((MODES, D))
for k in range(MODES):
    a = 3.0 if k < 4 else -3.0
    b = 3.0 if (k % 4) < 2 else -3.0
    base_means[k, :4] = [a, b, 2.2 * (1 if k % 2 else -1), 1.5 * (1 if k in (0, 3, 4, 7) else -1)]
# A shared target shift gives a clean, controllable nonstationarity signal.
shifts = []
for i in range(28):
    t = i / 27.0
    speed = 0.15 if i < 7 or i >= 21 else 0.95
    # piecewise rapid motion, plus a return, so consecutive q drift is observable
    x = (0.0 if i == 0 else sum(0.15 if j < 7 or j >= 21 else 0.95 for j in range(1, i+1)))
    shifts.append(np.array([x, -0.7*x, 0.35*x] + [0.] * (D-3)))
target_means = np.array([base_means + s for s in shifts])

def log_mix(x, means):
    # equal-weight spherical Gaussian mixture
    z = -0.5 * (((x[:, None, :] - means[None, :, :]) / SIGMA) ** 2).sum(axis=2)
    z -= D * (math.log(SIGMA) + 0.5 * LOG2PI)
    return logsumexp(z, axis=1) - math.log(MODES)

def sample_mix(means, n, rg):
    idx = rg.integers(MODES, size=n)
    return means[idx] + SIGMA * rg.normal(size=(n, D))

def ess_from_logw(logw):
    a = logw - logsumexp(logw)
    return float(np.exp(-logsumexp(2*a)) / len(logw))

def sym_kl_mc(a, b, n=5000):
    xa, xb = sample_mix(a, n, rng), sample_mix(b, n, rng)
    return float(0.5 * (np.mean(log_mix(xa, a)-log_mix(xa, b)) +
                        np.mean(log_mix(xb, b)-log_mix(xb, a))))

def importance_stats(target, proposal, n=M):
    x = sample_mix(proposal, n, rng)
    lw = log_mix(x, target) - log_mix(x, proposal)
    return ess_from_logw(lw), float(abs(np.log(np.mean(np.exp(lw)))))

def verify_math():
    a = base_means.copy()
    b = a + np.array([1.4, -0.5, 0.2] + [0.]*(D-3))
    # ESS is invariant to a common log-weight shift and lies in [1/M,1].
    lw = rng.normal(size=1000)
    e1 = ess_from_logw(lw)
    e2 = ess_from_logw(lw + 37.0)
    kl = sym_kl_mc(a, b, 12000)
    kl_rev = sym_kl_mc(b, a, 12000)
    # For this symmetric definition, swapping distributions should agree up to MC noise.
    return {"ess_in_range": bool(1/1000 - 1e-12 <= e1 <= 1.0),
            "ess_shift_invariance_abs_error": abs(e1-e2),
            "symmetric_kl_forward_reverse_abs_error": abs(kl-kl_rev),
            "symmetric_kl_positive": kl > 0, "symmetric_kl": kl}

def run(adaptive):
    q = base_means.copy()
    previous_q = q.copy()
    K, R = 8, 1
    records = []
    # persistent FIFO of target means emulates stale flow-training data.
    history = []
    for i in range(len(target_means)):
        history.append(target_means[i].copy())
        if len(history) > K: history.pop(0)
        # A gradient-like flow update toward the FIFO training distribution.
        train_target = np.mean(history, axis=0)
        for _ in range(int(R)):
            q += 0.30 * (train_target - q)
        drift = sym_kl_mc(q, previous_q, 1800) if i else 0.0
        eta, evidence_error = importance_stats(target_means[i], q)
        # Small local Langevin correction is represented by a conservative mode-center
        # pull only when ESS is poor and drift is not high; it is counted separately.
        correction = 0
        if adaptive and eta < 0.10 and drift <= 0.45:
            correction = 6
            q += 0.35 * (target_means[i] - q)
            eta, evidence_error = importance_stats(target_means[i], q)
        records.append({"level": i, "eta": eta, "drift": drift, "K": K, "R": R,
                        "correction_steps": correction, "log_evidence_abs_error": evidence_error})
        if adaptive:
            if drift > 0.45:
                K = max(2, K // 2)
            elif drift < 0.08 and eta > 0.45:
                K = min(8, K + 1)
            if eta < 0.10:
                R = min(8, max(1, int(math.ceil(R * (1 + 2.0*(0.10-eta))))))
            elif eta > 0.55:
                R = max(1, R-1)
        previous_q = q.copy()
    return records

def summarize(rs):
    eta = np.array([r["eta"] for r in rs])
    return {"mean_eta": float(eta.mean()), "min_eta": float(eta.min()),
            "levels_eta_below_0.1": int((eta < .1).sum()),
            "levels_eta_below_0.02": int((eta < .02).sum()),
            "mean_drift": float(np.mean([r["drift"] for r in rs])),
            "total_R": int(sum(r["R"] for r in rs)),
            "total_correction_steps": int(sum(r["correction_steps"] for r in rs)),
            "mean_log_evidence_abs_error": float(np.mean([r["log_evidence_abs_error"] for r in rs])),
            "max_log_evidence_abs_error": float(max(r["log_evidence_abs_error"] for r in rs))}

def main():
    mathcheck = verify_math()
    fixed, adaptive = run(False), run(True)
    out = {"seed": SEED, "math_check": mathcheck,
           "fixed_summary": summarize(fixed), "adaptive_summary": summarize(adaptive),
           "fixed_levels": fixed, "adaptive_levels": adaptive}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"math_check": mathcheck, "fixed": summarize(fixed), "adaptive": summarize(adaptive)}, indent=2))

if __name__ == "__main__": main()