import json
import numpy as np

SEED = 1037


def route(u, lam):
    return np.argmax(u - lam[None, :], axis=1)


def dual_bound(u, lam, cap):
    return float(lam @ cap + np.max(u - lam[None, :], axis=1).sum())


def accepted_stats(u, a, cap):
    counts = np.bincount(a, minlength=u.shape[1])
    # Capacity rule: keep highest-utility tokens within each expert.
    keep = np.zeros(len(a), dtype=bool)
    for e in range(u.shape[1]):
        ids = np.flatnonzero(a == e)
        if len(ids):
            ids = ids[np.argsort(u[ids, e])[::-1][:int(cap[e])]]
            keep[ids] = True
    p = float(u[np.arange(len(a))[keep], a[keep]].sum())
    return counts, keep, p


def math_checks(rng):
    # Prediction 1: weak duality has no violations for any nonnegative prices.
    violations = []
    gaps = []
    for _ in range(1000):
        n, e = 37, 5
        u = rng.uniform(0.05, 2.0, (n, e))
        cap = np.full(e, n / e)
        lam = rng.exponential(1.0, e)
        a = route(u, lam)
        _, keep, p = accepted_stats(u, a, cap)
        L = dual_bound(u, lam, cap)
        gaps.append(L - p)
        violations.append(max(0.0, p - L))
    # Prediction 2: local stability boundary for iid two-expert differences is rho*n < 2.
    # Around equal prices, load difference has slope approximately -n*(price difference),
    # so the linearized price-difference multiplier is 1-rho*n.
    n = 200
    rhos = [0.001, 0.005, 0.01, 0.02, 0.04]
    dyn = []
    for rho in rhos:
        lam = np.zeros(2)
        abs_imb = []
        for t in range(500):
            d = rng.uniform(-1, 1, n)
            u = np.column_stack([d, np.zeros(n)]) + 2.0
            a = route(u, lam)
            counts = np.bincount(a, minlength=2)
            lam = np.maximum(0, lam + rho * (counts - n / 2))
            if t >= 250:
                abs_imb.append(abs(counts[0] - counts[1]))
        dyn.append({"rho": rho, "predicted_multiplier": 1-rho*n,
                    "mean_abs_load_imbalance": float(np.mean(abs_imb)),
                    "max_abs_load_imbalance": int(np.max(abs_imb)),
                    "final_price_difference": float(lam[0]-lam[1])})
    # Prediction 3: at zero prices, routing is utility greedy; positive prices emerge
    # exactly when an expert is persistently over capacity.
    cap = np.array([50., 50.])
    n = 100
    u = np.zeros((n, 2)); u[:70, 0] = 1.0; u[70:, 1] = 1.0
    lam = np.zeros(2)
    zero_counts = np.bincount(route(u, lam), minlength=2)
    for _ in range(30):
        a = route(u, lam); c = np.bincount(a, minlength=2)
        lam = np.maximum(0, lam + .02*(c-cap))
    return {"weak_duality": {"trials":1000, "max_violation":float(max(violations)),
                             "min_gap":float(min(gaps)), "mean_gap":float(np.mean(gaps))},
            "stability_sweep": dyn,
            "zero_price_prediction": {"initial_greedy_counts":zero_counts.tolist(),
                                       "final_prices":lam.tolist(),
                                       "final_counts":np.bincount(route(u,lam),minlength=2).tolist(),
                                       "predicted": "overloaded expert gets positive price and its routed load falls"}}


def mini_experiment(rng):
    E, n, cap_each, batches = 4, 128, 32, 500
    bias = np.array([0.75, 0.35, 0.05, -0.2])
    def batch():
        return rng.normal(0, 0.35, (n,E)) + bias[None,:]
    results = {}
    for name, rho in [("baseline_greedy", None), ("dual_rho_0.001", .001),
                      ("dual_rho_0.01", .01), ("dual_rho_0.1", .1)]:
        lam = np.zeros(E); vals=[]; over=[]; vars_=[]; gaps=[]; prices=[]
        for t in range(batches):
            u = batch()
            a = route(u, np.zeros(E) if rho is None else lam)
            counts, keep, p = accepted_stats(u, a, np.full(E, cap_each))
            vals.append(p); over.append(int(np.maximum(counts-cap_each,0).sum()))
            vars_.append(float(np.var(counts)))
            if rho is not None:
                gaps.append(dual_bound(u,lam,np.full(E,cap_each))-p)
                lam = np.maximum(0, lam + rho*(counts-cap_each))
                prices.append(lam.copy())
        results[name] = {"mean_accepted_utility":float(np.mean(vals)),
                         "mean_overflow_tokens":float(np.mean(over)),
                         "mean_load_variance":float(np.mean(vars_)),
                         "final_prices":(lam.tolist() if rho is not None else None),
                         "mean_dual_gap":(float(np.mean(gaps)) if gaps else None),
                         "p95_dual_gap":(float(np.percentile(gaps,95)) if gaps else None)}
    return results


def main():
    rng = np.random.default_rng(SEED)
    out = {"seed":SEED, "math":math_checks(rng), "mini_experiment":mini_experiment(rng)}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()