import json
import math
import numpy as np


def dg_moments(q, limit=200):
    z = np.arange(-limit, limit + 1)
    p = (1-q)/(1+q) * q**np.abs(z)
    p /= p.sum()
    t = np.abs(z)
    return p, z, float((p*t).sum()), float((p*t*t).sum() - (p*t).sum()**2)


def math_check():
    q = 0.37
    p, z, mu_emp, var_emp = dg_moments(q)
    mu = 2*q/(1-q*q)
    # Exact expectation and derivative for a quadratic objective at one coordinate.
    target, x = 2, -1
    def expected(eta):
        qq = math.exp(eta)
        pp = (1-qq)/(1+qq) * qq**np.abs(z)
        return float((pp * (x + z - target)**2).sum())
    eta = math.log(q)
    eps = 2e-5
    fd = (expected(eta+eps)-expected(eta-eps))/(2*eps)
    F = (x + z - target)**2
    score_cov = float((p*(F-(p*F).sum())*(np.abs(z)-mu)).sum())
    fisher = var_emp
    return {
        "q": q, "pmf_sum": float(p.sum()), "analytic_mu": mu,
        "empirical_mu": mu_emp, "empirical_fisher": fisher,
        "score_identity_cov": score_cov, "finite_difference_dE_deta": fd,
        "identity_abs_error": abs(score_cov-fd),
        "natural_gradient": score_cov/(fisher+1e-12),
    }


def utilities(losses):
    # Higher utility means better (lower) loss; centered rank utilities.
    ranks = np.argsort(np.argsort(losses))
    u = (len(losses)-1 - 2*ranks).astype(float) / max(1, len(losses)-1)
    return u - u.mean()


def dg_mutation(rng, q, d):
    mag = rng.geometric(1-q, size=d) - 1
    sign = np.where(rng.random(d) < .5, -1, 1)
    return mag * sign


def optimize(method, seed, targets, weights, iterations=90, K=24):
    rng = np.random.default_rng(seed)
    d = len(targets)
    x = np.zeros(d, dtype=int)
    q = np.full(d, .35)
    eta = np.log(q)
    h = np.zeros(d)
    beta, rho = .85, .22
    best = float(np.sum(weights*(x-targets)**2))
    curve = []
    for _ in range(iterations):
        zs = []
        if method == "double_geometric":
            for _k in range(K): zs.append(dg_mutation(rng, q, d))
        elif method == "gaussian_es":
            # Standard Gaussian ES, rounded to the same integer search space.
            zs = [np.rint(rng.normal(0, 1.5, size=d)).astype(int) for _k in range(K)]
        elif method == "fixed_radius":
            # Simple integer random search with independent {-1,0,+1} mutations.
            zs = [rng.integers(-1, 2, size=d) for _k in range(K)]
        losses = np.array([np.sum(weights*(np.clip(x+z, -10, 10)-targets)**2) for z in zs])
        i = int(np.argmin(losses))
        if losses[i] < best: best = float(losses[i])
        curve.append(best)
        # All methods use the same rank-based selection signal; only DG adapts q.
        if method == "double_geometric":
            U = utilities(losses)
            T = np.array([np.abs(z) for z in zs], dtype=float)
            mu = 2*q/(1-q*q)
            var = T.var(axis=0)
            # Cov(utility,T)/Var(T) estimates negative objective natural gradient.
            g_utility = np.mean(U[:,None]*(T-mu), axis=0)/(var+1e-6)
            h = beta*h + (1-beta)*g_utility
            eta = np.clip(eta + rho*h, math.log(.08), math.log(.80))
            q = np.exp(eta)
        # ES update: select the best candidate, matching a minimal zero-order optimizer.
        x = np.clip(x + zs[i], -10, 10).astype(int)
    return {"final_best": best, "curve": curve, "final_x": x.tolist(), "final_q": q.tolist()}


def main():
    check = math_check()
    # Six layerwise integer controls with heterogeneous scales/targets.
    targets = np.array([4, -3, 2, 5, -2, 3])
    weights = np.array([1., 1.4, .8, 1.2, .9, 1.1])
    methods = ["double_geometric", "gaussian_es", "fixed_radius"]
    results = {m: [optimize(m, s, targets, weights) for s in range(8)] for m in methods}
    summary = {}
    for m, rr in results.items():
        finals = np.array([r["final_best"] for r in rr])
        # Equal evaluation budget: report final objective and fraction reaching exact target.
        summary[m] = {"final_mean": float(finals.mean()), "final_std": float(finals.std()),
                       "exact_fraction": float(np.mean(finals < 1e-9)),
                       "median": float(np.median(finals))}
    out = {"math_check": check, "summary": summary, "raw": results,
           "evaluations_per_run": 90*24, "targets": targets.tolist()}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps({"math_check": check, "summary": summary}, indent=2))

if __name__ == "__main__": main()