import json
import numpy as np
from scipy.optimize import minimize_scalar

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

# A bounded activation model makes the infinite-domain construction numerically
# concrete while preserving the paper's local asymptotics.
XMAX = 8.0
RATE = 1.0
SCALE = 0.5


def sample_x(n=400000):
    u = rng.random(n)
    # inverse CDF of Exp(1) conditional on x <= XMAX
    return -np.log(1.0 - u * (1.0 - np.exp(-RATE * XMAX))) / RATE


def density(x):
    z = 1.0 - np.exp(-RATE * XMAX)
    return RATE * np.exp(-RATE * x) / z


def utility(x):
    return np.log1p(np.maximum(x, 0.0) / SCALE) / np.log1p(XMAX / SCALE)


def utility_prime(x):
    return 1.0 / ((SCALE + x) * np.log1p(XMAX / SCALE))


def weighted_thresholds(n, grid_n=200001):
    grid = np.linspace(0.0, XMAX, grid_n)
    g = np.sqrt(density(grid) * utility_prime(grid))
    # cumulative trapezoidal integration
    cdf = np.concatenate([[0.0], np.cumsum((g[:-1] + g[1:]) * np.diff(grid) * 0.5)])
    cdf /= cdf[-1]
    targets = np.arange(1, n) / n
    internal = np.interp(targets, cdf, grid)
    return np.r_[0.0, internal, XMAX]


def uniform_thresholds(n):
    return np.linspace(0.0, XMAX, n + 1)


def left_edge_quantize(x, edges):
    # edges has n+1 entries; the last edge is the finite support endpoint.
    idx = np.searchsorted(edges, x, side="right") - 1
    return edges[np.clip(idx, 0, len(edges) - 2)]


def lloyd_max(x, n, iterations=80):
    # Standard MSE scalar quantizer, with nearest-centroid reconstruction.
    cent = np.quantile(x, (np.arange(n) + 0.5) / n)
    for _ in range(iterations):
        boundaries = (cent[:-1] + cent[1:]) * 0.5
        labels = np.searchsorted(boundaries, x)
        new = cent.copy()
        for k in range(n):
            vals = x[labels == k]
            if len(vals):
                new[k] = vals.mean()
        if np.max(np.abs(new - cent)) < 1e-10:
            break
        cent = new
    return cent


def lloyd_quantize(x, cent):
    boundaries = (cent[:-1] + cent[1:]) * 0.5
    return cent[np.searchsorted(boundaries, x)]


def utility_constant():
    grid = np.linspace(0.0, XMAX, 400001)
    g = np.sqrt(density(grid) * utility_prime(grid))
    z = np.trapz(g, grid)
    return 0.5 * z * z


def evaluate(x, q, ideal_utility=None):
    if ideal_utility is None:
        ideal_utility = utility(x)
    gap = ideal_utility - utility(q)
    return {
        "utility_gap": float(np.mean(gap)),
        "n_gap": None,
        "mse": float(np.mean((x - q) ** 2)),
        "upward_fraction": float(np.mean(q > x + 1e-12)),
        "mean_abs_error": float(np.mean(np.abs(x - q))),
    }


def main():
    x = sample_x()
    ideal = utility(x)
    c_theory = utility_constant()
    rows = []
    scaling = []
    for n in [4, 8, 16, 32, 64]:
        ew = weighted_thresholds(n)
        eu = uniform_thresholds(n)
        qw = left_edge_quantize(x, ew)
        qu = left_edge_quantize(x, eu)
        cent = lloyd_max(x, n)
        ql = lloyd_quantize(x, cent)
        rw = evaluate(x, qw, ideal)
        ru = evaluate(x, qu, ideal)
        rl = evaluate(x, ql, ideal)
        rw["n_gap"] = n * rw["utility_gap"]
        ru["n_gap"] = n * ru["utility_gap"]
        rl["n_gap"] = n * rl["utility_gap"]
        rows.append({"n": n, "weighted_left": rw, "uniform_left": ru, "mse_lloyd": rl})
        scaling.append({"n": n, "weighted_n_gap": rw["n_gap"], "uniform_n_gap": ru["n_gap"], "theory_Cw": c_theory})

    # Direct numerical check of the local allocation rule: weighted bins have
    # nearly equal integrals of sqrt(p Q'), unlike uniform bins.
    e = weighted_thresholds(32)
    grid = np.linspace(0, XMAX, 300001)
    g = np.sqrt(density(grid) * utility_prime(grid))
    gcdf = np.concatenate([[0], np.cumsum((g[:-1] + g[1:]) * np.diff(grid) / 2)])
    bin_mass = np.diff(np.interp(e, grid, gcdf))
    allocation_check = {
        "weighted_g_mass_cv": float(np.std(bin_mass) / np.mean(bin_mass)),
        "weighted_g_mass_min": float(bin_mass.min()),
        "weighted_g_mass_max": float(bin_mass.max()),
    }

    result = {
        "seed": SEED,
        "distribution": "Exp(rate=1) truncated to [0,8]",
        "utility": "log1p(x/0.5)/log1p(8/0.5)",
        "theory_Cw": c_theory,
        "results": rows,
        "scaling": scaling,
        "allocation_check": allocation_check,
        "interpretation": {
            "claimed_signal": "weighted left-edge utility gap scales approximately as Cw/n and is lower than simple uniform left-edge",
            "conservative_rule": "weighted and uniform left-edge have zero upward rounding by construction",
            "mse_control": "Lloyd-Max minimizes squared error, not the monotone left-edge utility gap",
        },
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()