import json
from pathlib import Path
import numpy as np


def softmax(x):
    x = x - np.max(x)
    p = np.exp(x)
    return p / p.sum()


def weighted_moments(z, w):
    mu = np.sum(w[:, None] * z, axis=0)
    c = np.sum(w[:, None] * (z - mu) ** 2, axis=0)
    return mu, c


def curvature_estimate(z, w, rewards, eps=1e-10):
    mu, c = weighted_moments(z, w)
    q = (z - mu) ** 2 - c
    rbar = np.sum(w * rewards)
    # This is the weighted least-squares coefficient in Cov(q,r)/E[q^2].
    r2 = np.sum(w[:, None] * (rewards[:, None] - rbar) * q, axis=0)
    den = np.sum(w[:, None] * q * q, axis=0) + eps
    return r2 / den, mu, c, rbar


def verify_moments(seed=7, n=2000000):
    rng = np.random.default_rng(seed)
    mu, s, kappa, curvature, slope = 0.7, 0.8, 1.3, -1.7, 0.4
    x = rng.normal(mu, s, n)
    r = slope * (x - mu) + 0.5 * curvature * (x - mu) ** 2
    cov_xr = np.mean((x - np.mean(x)) * (r - np.mean(r)))
    q = (x - np.mean(x)) ** 2 - np.var(x)
    cov_qr = np.mean((q - np.mean(q)) * (r - np.mean(r)))
    mean_rhs = kappa * s**2 * slope
    var_rhs_selection = kappa * curvature * s**4
    # Independent diffusion check: Gaussian perturbation variance increases by 2D dt.
    D, dt = 0.23, 0.01
    x2 = x + np.sqrt(2 * D * dt) * rng.normal(size=n)
    diffusion_increment = np.var(x2) - np.var(x)
    return {
        "empirical_mean_derivative_rhs": float(kappa * cov_xr),
        "predicted_mean_derivative": float(mean_rhs),
        "empirical_variance_selection_rhs": float(kappa * cov_qr),
        "predicted_variance_selection": float(var_rhs_selection),
        "empirical_diffusion_variance_increment": float(diffusion_increment),
        "predicted_diffusion_variance_increment": float(2 * D * dt),
        "relative_mean_error": float(abs(kappa * cov_xr - mean_rhs) / (abs(mean_rhs) + 1e-12)),
        "relative_variance_error": float(abs(kappa * cov_qr - var_rhs_selection) / (abs(var_rhs_selection) + 1e-12)),
    }


def run_population(mode, seed=123, steps=500, n=16, d=8):
    rng = np.random.default_rng(seed)
    target = np.linspace(-1.0, 1.0, d)
    z = rng.normal(0.0, 1.0, (n, d))
    # Fixed mutation is deliberately a reasonable, nonzero control.
    fixed_D = 0.018
    dmin, dmax, gamma, target_c = 0.00015, 0.045, 0.20, 0.16
    kappa, dt = 2.0, 1.0
    records = []
    for t in range(steps):
        true_reward = -np.mean((z - target) ** 2, axis=1)
        noise = rng.normal(0, 0.10, n)
        rewards = true_reward + noise
        uniform = np.full(n, 1.0 / n)
        r2, mu, c, rbar = curvature_estimate(z, uniform, rewards)
        reward_var = np.var(rewards)
        # Adaptive diffusion follows the proposal: noisy rewards increase D;
        # only sufficiently reliable negative curvature permits reduction.
        Dcoord = np.clip(dmin + gamma * reward_var / (1.0 + np.abs(np.clip(r2, -8, 8))), dmin, dmax)
        ess = 1.0 / np.sum(softmax(kappa * (rewards - rbar)) ** 2)
        if mode == "fixed":
            Dcoord = np.full(d, fixed_D)
        elif mode == "moment":
            if np.all(r2 < 0) and ess > n / 2:
                Dcoord = np.maximum(Dcoord, dmin)
            # Explicit +2D safeguard for a coordinate below target diversity.
            Dcoord = np.maximum(Dcoord, np.maximum(0.0, (target_c - c) / (2 * dt)))
            Dcoord = np.clip(Dcoord, dmin, dmax)
        weights = softmax(kappa * (rewards - rbar))
        # Resample around the weighted selected center, as in a population adapter optimizer.
        parent_idx = rng.choice(n, size=n, p=weights)
        z = z[parent_idx] + np.sqrt(2.0 * Dcoord * dt)[None, :] * rng.normal(size=(n, d))
        if t % 10 == 0 or t == steps - 1:
            records.append({
                "step": t, "best_true_reward": float(np.max(true_reward)),
                "mean_true_reward": float(np.mean(true_reward)),
                "distance": float(np.mean((np.mean(z, axis=0) - target) ** 2)),
                "cov_trace": float(np.mean(np.var(z, axis=0))), "ess": float(ess),
                "reward_var": float(reward_var), "mutation_D": float(np.mean(Dcoord)),
            })
    return records


def main():
    out = Path("results")
    out.mkdir(exist_ok=True)
    verification = verify_moments()
    all_runs = {m: run_population(m) for m in ("fixed", "moment")}
    summary = {"verification": verification, "runs": all_runs}
    for mode, rows in all_runs.items():
        with (out / (mode + ".jsonl")).open("w") as f:
            for row in rows:
                f.write(json.dumps(row) + "\n")
    with (out / "summary.json").open("w") as f:
        json.dump(summary, f, indent=2)
    for mode, rows in all_runs.items():
        a, b = rows[0], rows[-1]
        print(mode, "final_best=%.4f final_mean=%.4f final_cov=%.4f final_D=%.4f" %
              (b["best_true_reward"], b["mean_true_reward"], b["cov_trace"], b["mutation_D"]))
    print("verification", json.dumps(verification))

if __name__ == "__main__":
    main()