"""Claim 5: Lemma 2.1 (argmax generically a singleton) and Lemma 2.2 (the convex
hull of tokens shrinks under step sizes gamma in (0,1)), underpinning the
Frank-Wolfe reformulation of hardmax attention.

Previously unaddressed: neither lemma was tested.

Hardmax attention as a Frank-Wolfe step: each token moves toward the token that
maximises its inner product,
    j*(i) = argmax_j <x_i, x_j>,      x_i <- (1-gamma) x_i + gamma x_{j*(i)}
which is exactly a Frank-Wolfe step toward a VERTEX of the current convex hull.
"""
import json, numpy as np
from scipy.spatial import ConvexHull

RES = {}


def step(X, gamma):
    S = X @ X.T
    order = np.argsort(-S, axis=1)
    j = order[:, 0]
    top2 = S[np.arange(len(X)), order[:, 0]]-S[np.arange(len(X)), order[:, 1]]
    return (1-gamma)*X+gamma*X[j], top2


def run():
    # ---- Lemma 2.1: argmax is generically a singleton ----
    rows = []
    for d in (2, 3, 5, 10):
        for n in (20, 50):
            gaps, ties = [], 0; tot = 0
            for s in range(40):
                rng = np.random.default_rng(s)
                X = rng.normal(size=(n, d))
                _, g = step(X, 0.3)
                gaps.append(float(np.min(g))); ties += int(np.sum(g <= 1e-12)); tot += n
            rows.append({"d": d, "n": n, "trials": 40, "tokens_checked": tot,
                         "exact_ties": ties, "min_argmax_gap": min(gaps)})
            print("  L2.1 d=%-3d n=%-3d  exact ties %d/%d   smallest argmax gap = %.3e"
                  % (d, n, ties, tot, min(gaps)), flush=True)
    RES["lemma21_singleton"] = {"rows": rows,
        "total_tokens": sum(r["tokens_checked"] for r in rows),
        "total_exact_ties": sum(r["exact_ties"] for r in rows),
        "smallest_gap_overall": min(r["min_argmax_gap"] for r in rows)}

    # ---- Lemma 2.2: convex hull shrinks for gamma in (0,1) ----
    rows = []
    for gamma in (0.1, 0.3, 0.5, 0.9):
        for d in (2, 3):
            vols, diams, mono_v, mono_d, runs = [], [], 0, 0, 0
            for s in range(12):
                rng = np.random.default_rng(100+s)
                X = rng.normal(size=(40, d))
                V, D = [], []
                for t in range(25):
                    try: V.append(ConvexHull(X).volume)
                    except Exception: V.append(0.0)
                    D.append(float(np.max(np.linalg.norm(X[:, None]-X[None, :], axis=-1))))
                    X, _ = step(X, gamma)
                V = np.array(V); D = np.array(D)
                mono_v += int(np.all(np.diff(V) <= 1e-9)); mono_d += int(np.all(np.diff(D) <= 1e-9))
                vols.append(float(V[-1]/max(V[0], 1e-12))); diams.append(float(D[-1]/max(D[0], 1e-12)))
                runs += 1
            rows.append({"gamma": gamma, "d": d, "runs": runs,
                         "volume_monotone_runs": mono_v, "diameter_monotone_runs": mono_d,
                         "final_over_initial_volume": round(float(np.mean(vols)), 6),
                         "final_over_initial_diameter": round(float(np.mean(diams)), 6)})
            print("  L2.2 gamma=%.1f d=%d  volume non-increasing %d/%d, diameter %d/%d | vol ratio %.4f, diam ratio %.4f"
                  % (gamma, d, mono_v, runs, mono_d, runs,
                     rows[-1]["final_over_initial_volume"], rows[-1]["final_over_initial_diameter"]), flush=True)
    RES["lemma22_hull_shrinks"] = {"rows": rows,
        "all_volume_monotone": all(r["volume_monotone_runs"] == r["runs"] for r in rows),
        "all_diameter_monotone": all(r["diameter_monotone_runs"] == r["runs"] for r in rows)}
    json.dump(RES, open("fw_results.json", "w"), indent=1)


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