#!/usr/bin/env python3
"""Fresh CPU audit of the Theorem 5.2 upper envelope.

This intentionally tests the upper-bound statement rather than fitting an
asymptotic n exponent.  For each family and (n,k) setting it computes the
Lipschitz constant of the optimized welfare by central finite differences on
fresh utility vectors, then runs the unchanged SWF-UCB simulator at three
horizons and four seeds.  The reported statistic is the empirical 90th
percentile of regret divided by L*(n+sqrt(n*k*T)).
"""

from __future__ import annotations

import json
import math
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from reproduce import oracle, simulate, welfare


FAMILIES = (("wpm", -1.0), ("kolm", -1.0), ("gini", 0.0))
NS = (8, 16, 32)
HORIZONS = (1024, 4096, 16384)
SEEDS = 4


def optimized_value(family, param, u, w, k):
    p, _ = oracle(family, u, w, param, k)
    return welfare(family, u, p, w, param)


def numerical_lipschitz(family, param, n, k, rng):
    """Executed max coordinate slope of M(u)=max_p M(u*p)."""
    maxima = []
    h = 1e-6
    for _ in range(12):
        u = rng.uniform(0.18, 0.92, n)
        w = np.full(n, 1.0 / n) if family != "gini" else np.linspace(2.0, 0.5, n)
        w = w / w.sum()
        for i in range(n):
            lo = u[i] - h
            hi = u[i] + h
            um = u.copy(); up = u.copy()
            um[i] = lo; up[i] = hi
            derivative = (optimized_value(family, param, up, w, k)
                          - optimized_value(family, param, um, w, k)) / (2.0 * h)
            maxima.append(abs(float(derivative)))
    return float(max(maxima))


def main(output: str):
    rows = []
    rng = np.random.default_rng(20260729)
    for fidx, (family, param) in enumerate(FAMILIES):
        for n in NS:
            weights = np.full(n, 1.0 / n) if family != "gini" else np.linspace(2.0, 0.5, n)
            weights = weights / weights.sum()
            for k in sorted(set((1, max(2, n // 4), max(2, n // 2)))):
                if k > n:
                    continue
                means = np.linspace(0.20, 0.90, n)
                means = means[rng.permutation(n)]
                local_L = numerical_lipschitz(family, param, n, k, rng)
                # Conservative global constants for the actual normalized
                # welfare families used here: for uniform-weight WPM(q=-1),
                # |dM/du_i| <= 1/w_i = n; Kolm and weighted sorted Gini are
                # 1-Lipschitz in ||.||_infty because their weights sum to one.
                # Keep the finite-difference slope separately as a diagnostic.
                L = float(n if family == "wpm" else 1.0)
                for horizon in HORIZONS:
                    regrets = [simulate(family, param, means, weights, k, horizon,
                                         940000 + 100000 * fidx + 1000 * n + 100 * k
                                         + horizon + seed)
                               for seed in range(SEEDS)]
                    q90 = float(np.quantile(regrets, 0.90, method="linear"))
                    envelope = L * (n + math.sqrt(n * k * horizon))
                    rows.append({
                        "family": family, "n": n, "k": k, "T": horizon,
                        "seeds": SEEDS, "L": L, "local_L": local_L,
                        "regret_q90": q90,
                        "envelope": envelope,
                        "q90_over_envelope": q90 / envelope,
                        "regret_mean": float(np.mean(regrets)),
                    })
                    print(rows[-1], flush=True)
    ratios = [r["q90_over_envelope"] for r in rows]
    summary = {
        "cells": len(rows), "families": [f for f, _ in FAMILIES],
        "n_values": list(NS), "horizons": list(HORIZONS), "seeds_per_cell": SEEDS,
        "max_q90_over_L_n_plus_sqrt_nkT": max(ratios),
        "p99_q90_over_envelope": float(np.quantile(ratios, 0.99)),
        "all_cells_below_constant_1": bool(max(ratios) <= 1.0),
        "fitted_smallest_uniform_constant": max(ratios),
    }
    result = {"protocol": "Theorem 5.2 CPU upper-envelope audit",
              "summary": summary, "rows": rows}
    Path(output).write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2), flush=True)


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "/tmp/social-claim2-upper-bound.json")