"""Claim 4: on 400x800 random zero-sum games, second-order (PSSN-style) methods
reach a duality gap of ~1e-12 in a few seconds.

The logbook checked this only as a source-reported figure. The claim has two
measurable parts -- the ACHIEVABLE GAP and the TIME -- and a second-order method
should separate sharply from a first-order one on the first of these.

Zero-sum game: min_x max_y x^T A y over the simplices. Duality gap at (x,y):
    gap = max_j (A^T x)_j  -  min_i (A y)_i .
Compared: extragradient / mirror-prox (first order) against an interior-point
solve of the equivalent LP (second order, Newton steps on the KKT system).
"""
import json, time, numpy as np
from scipy.optimize import linprog

RES = {}


def game(m, n, seed):
    return np.random.default_rng(seed).normal(size=(m, n))


def gap(A, x, y):
    return float((A.T @ x).max()-(A @ y).min())


def simplex_proj(v):
    u = np.sort(v)[::-1]; c = np.cumsum(u)-1
    r = np.arange(1, len(v)+1)
    rho = np.nonzero(u-c/r > 0)[0][-1]
    return np.maximum(v-c[rho]/(rho+1), 0)


def first_order(A, iters=20000):
    m, n = A.shape
    x = np.ones(m)/m; y = np.ones(n)/n
    L = np.linalg.norm(A, 2); eta = 0.9/max(L, 1e-12)
    t0 = time.perf_counter()
    for _ in range(iters):
        xh = simplex_proj(x-eta*(A @ y)); yh = simplex_proj(y+eta*(A.T @ x))
        x = simplex_proj(x-eta*(A @ yh)); y = simplex_proj(y+eta*(A.T @ xh))
    return gap(A, x, y), time.perf_counter()-t0


def second_order(A):
    """LP form solved by interior point (Newton on the KKT system)."""
    m, n = A.shape
    t0 = time.perf_counter()
    # min_x max_j (A^T x)_j  ->  min t s.t. A^T x <= t 1, sum x = 1, x >= 0
    c = np.zeros(m+1); c[-1] = 1.0
    Aub = np.hstack([A.T, -np.ones((n, 1))])
    Aeq = np.zeros((1, m+1)); Aeq[0, :m] = 1.0
    r = linprog(c, A_ub=Aub, b_ub=np.zeros(n), A_eq=Aeq, b_eq=[1.0],
                bounds=[(0, None)]*m+[(None, None)], method="highs-ipm")
    x = r.x[:m]
    c2 = np.zeros(n+1); c2[-1] = -1.0
    Aub2 = np.hstack([-A, np.ones((m, 1))])
    Aeq2 = np.zeros((1, n+1)); Aeq2[0, :n] = 1.0
    r2 = linprog(c2, A_ub=Aub2, b_ub=np.zeros(m), A_eq=Aeq2, b_eq=[1.0],
                 bounds=[(0, None)]*n+[(None, None)], method="highs-ipm")
    y = r2.x[:n]
    el = time.perf_counter()-t0
    return gap(A, x, y), el


def main():
    rows = []
    for (m, n) in ((100, 200), (200, 400), (400, 800)):
        for s in range(3):
            A = game(m, n, seed=s)
            g1, t1 = first_order(A)
            g2, t2 = second_order(A)
            rows.append({"m": m, "n": n, "seed": s,
                         "first_order_gap": g1, "first_order_sec": round(t1, 3),
                         "second_order_gap": g2, "second_order_sec": round(t2, 3)})
        sub = [r for r in rows if r["m"] == m]
        print("  %dx%d  first-order: gap=%.3e in %.2fs   |   second-order: gap=%.3e in %.2fs"
              % (m, n, np.mean([r["first_order_gap"] for r in sub]), np.mean([r["first_order_sec"] for r in sub]),
                 np.mean([r["second_order_gap"] for r in sub]), np.mean([r["second_order_sec"] for r in sub])), flush=True)
    big = [r for r in rows if r["m"] == 400]
    RES["claim4_pssn"] = {"rows": rows, "seeds": 3,
        "target_size": "400x800",
        "second_order_gap_400x800": float(np.mean([r["second_order_gap"] for r in big])),
        "second_order_sec_400x800": round(float(np.mean([r["second_order_sec"] for r in big])), 3),
        "first_order_gap_400x800": float(np.mean([r["first_order_gap"] for r in big])),
        "gap_ratio": float(np.mean([r["first_order_gap"] for r in big])/max(np.mean([r["second_order_gap"] for r in big]), 1e-300)),
        "paper_reported": {"gap": 1e-12, "seconds": "4.3-4.4"}}
    R = RES["claim4_pssn"]
    print("  400x800: second-order gap %.3e in %.2fs (paper: 1e-12 in 4.3-4.4s); first-order gap %.3e -> %.1e x worse"
          % (R["second_order_gap_400x800"], R["second_order_sec_400x800"],
             R["first_order_gap_400x800"], R["gap_ratio"]), flush=True)
    json.dump(RES, open("pssn_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")