import json
import math
import numpy as np


def project_simplex(v, z=1.0):
    """Euclidean projection onto {x >= 0, sum(x)=z}."""
    v = np.asarray(v, dtype=float)
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u) - z
    rho = np.nonzero(u - cssv / (np.arange(len(v)) + 1) > 0)[0]
    if len(rho) == 0:
        return np.full_like(v, z / len(v))
    theta = cssv[rho[-1]] / (rho[-1] + 1)
    return np.maximum(v - theta, 0.0)


def payoff(x, a, b):
    # x: groups x experts; a: groups x experts; b: experts
    D = b + x.sum(axis=0)
    return float(np.sum(a * x / D[None, :]))


def equimarginal_route(a, b, budget=1.0, steps=20, eta=0.5,
                       tol=1e-10, backtrack=True, init=None):
    """Simultaneous projected ascent on each group's own payoff.

    The update is x_j <- simplex_project(x_j + eta*g_j), where
    g_ji = a_ji*(D_i-x_ji)/D_i^2. Backtracking keeps batch payoff monotone.
    """
    a = np.asarray(a, dtype=float)
    b = np.asarray(b, dtype=float)
    m, n = a.shape
    x = np.full((m, n), budget / n) if init is None else np.array(init, float, copy=True)
    history = [payoff(x, a, b)]
    for _ in range(steps):
        D = b + x.sum(axis=0)
        g = a * (D[None, :] - x) / (D[None, :] ** 2)
        local_eta = eta
        old = history[-1]
        while True:
            y = np.vstack([project_simplex(row + local_eta * grad, budget)
                           for row, grad in zip(x, g)])
            new = payoff(y, a, b)
            if (not backtrack) or new >= old - 1e-12 or local_eta < 1e-9:
                break
            local_eta *= 0.5
        x = y
        history.append(new)
        if np.max(np.abs(g - g.mean(axis=1, keepdims=True))) < tol:
            break
    return x, np.asarray(history)


def independent_softmax(a, temperature=1.0, budget=1.0):
    z = a / max(temperature, 1e-12)
    z -= z.max(axis=1, keepdims=True)
    p = np.exp(z)
    p /= p.sum(axis=1, keepdims=True)
    return budget * p


def cv(v):
    v = np.asarray(v)
    return float(v.std() / (v.mean() + 1e-12))


def transition_check():
    # One-player water filling has x_i=sqrt(a_i*b_i/c)-b_i and
    # a new expert enters when c=a_i/b_i. These are exact predictions.
    a = np.array([[4.0, 2.0, 1.0, 0.25]])
    b = np.ones(4)
    ratios = (a[0] / b)
    predicted = []
    for k in range(1, 4):
        c = ratios[k]
        r = np.sum(np.sqrt(a[0, :k] * b[:k] / c) - b[:k])
        predicted.append(float(r))
    # Probe just below/above each threshold and detect positive support.
    observed = []
    for r0 in predicted:
        for r in (r0 * (1 - 2e-3), r0 * (1 + 2e-3)):
            x, _ = equimarginal_route(a, b, budget=r, steps=500, eta=0.8)
            observed.append({"budget": float(r), "support": int((x[0] > 1e-6).sum())})
    return predicted, observed


def marginal_check():
    rng = np.random.default_rng(7)
    a = np.exp(rng.normal(0, .7, size=(5, 6)))
    b = np.exp(rng.normal(0, .3, size=6))
    x, hist = equimarginal_route(a, b, steps=80, eta=1.0)
    D = b + x.sum(axis=0)
    g = a * (D[None, :] - x) / D[None, :] ** 2
    residuals = []
    for j in range(len(a)):
        active = x[j] > 1e-5
        c = np.median(g[j, active]) if active.any() else 0
        # KKT residual: active rates should equal c; inactive rates <= c.
        active_err = np.max(np.abs(g[j, active] - c)) if active.any() else 0
        inactive_violation = np.max(np.maximum(g[j, ~active] - c, 0)) if (~active).any() else 0
        residuals.append(float(max(active_err, inactive_violation)))
    return float(hist[0]), float(hist[-1]), float(max(residuals)), x


def load_sweep():
    # Same affinity pattern for all groups: independent routing overloads the
    # high-affinity experts; congestion routing's predicted response is to
    # spread mass as group count raises congestion.
    rng = np.random.default_rng(11)
    n, mmax = 8, 64
    b = np.ones(n) * 0.25
    base = np.array([3.0, 2.6, 2.2, 1.8, 1.4, 1.1, .9, .7])
    rows = []
    for m in [4, 8, 16, 32, 64]:
        noise = rng.normal(0, .05, size=(m, n))
        a = np.maximum(base[None, :] * np.exp(noise), 1e-5)
        x_c, h = equimarginal_route(a, b, steps=30, eta=1.0)
        x_s = independent_softmax(np.log(a), temperature=.7)
        cap = (m / n) * 1.15
        rows.append({"groups": m,
                     "congestion_cv": cv(x_c.sum(0)),
                     "softmax_cv": cv(x_s.sum(0)),
                     "congestion_overflow": float(np.maximum(x_c.sum(0)-cap, 0).sum()),
                     "softmax_overflow": float(np.maximum(x_s.sum(0)-cap, 0).sum()),
                     "congestion_payoff": payoff(x_c, a, b),
                     "softmax_payoff": payoff(x_s, a, b),
                     "iterations": len(h)-1})
    return rows


def main():
    pred, obs = transition_check()
    initial, final, residual, x = marginal_check()
    loads = load_sweep()
    result = {"transition_predicted_budgets": pred,
              "transition_observed": obs,
              "marginal_initial_payoff": initial,
              "marginal_final_payoff": final,
              "marginal_kkt_residual": residual,
              "marginal_row_sum_max_error": float(np.max(np.abs(x.sum(1)-1))),
              "load_sweep": loads}
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()