#!/usr/bin/env python3
"""Timing audit for the released Theorem 5.1 policy-oracle implementations."""

from __future__ import annotations

import json
import math
import statistics
import sys
import time
from pathlib import Path

import numpy as np


ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "artifact"))
from core.oracles import oracle_gini, oracle_kolm, oracle_wpm  # noqa: E402


SIZES = [256, 512, 1024, 2048, 4096, 8192]
REPEATS = 11
SEED = 20260731


def slope(rows: list[dict]) -> float:
    x = np.log([row["n"] for row in rows])
    y = np.log([row["median_seconds"] for row in rows])
    return float(np.polyfit(x, y, 1)[0])


def timed(function, arguments) -> float:
    samples = []
    for _ in range(REPEATS):
        started = time.perf_counter()
        output = function(*arguments)
        samples.append(time.perf_counter() - started)
        assert np.all(np.isfinite(output))
        assert np.all(output >= -1e-12)
        assert np.all(output <= 1 + 1e-12)
        assert abs(float(output.sum()) - arguments[-1]) < 1e-8
    return statistics.median(samples)


def main() -> None:
    rng = np.random.default_rng(SEED)
    results = {"wpm": [], "kolm": [], "gini": []}
    for n in SIZES:
        k = min(16, n // 4)
        u = rng.uniform(0.1, 1.0, n)
        weights = rng.uniform(0.1, 1.0, n)
        weights /= weights.sum()
        cases = {
            "wpm": (oracle_wpm, (u, weights, -1.0, k)),
            "kolm": (oracle_kolm, (u, weights, -1.0, k)),
            "gini": (oracle_gini, (u, weights, k)),
        }
        for name, (function, arguments) in cases.items():
            seconds = timed(function, arguments)
            denominator = (
                n * math.log2(n)
                if name in {"wpm", "kolm"}
                else k * n
            )
            results[name].append(
                {
                    "n": n,
                    "k": k,
                    "median_seconds": seconds,
                    "seconds_per_theory_unit": seconds / denominator,
                }
            )

    summary = {
        "seed": SEED,
        "sizes": SIZES,
        "repeats_per_cell": REPEATS,
        "fixed_k": 16,
        "implementation": "JG1310/repro-swf-allocation-bundle core/oracles.py",
        "source_sha256": __import__("hashlib").sha256(
            (ROOT / "artifact" / "core" / "oracles.py").read_bytes()
        ).hexdigest(),
        "results": results,
        "log_log_slopes": {
            name: slope(rows) for name, rows in results.items()
        },
        "checks": {
            "all_outputs_feasible": True,
            "wpm_slope_below_1_5": slope(results["wpm"]) < 1.5,
            "kolm_slope_below_1_5": slope(results["kolm"]) < 1.5,
            "gini_fixed_k_slope_below_1_5": slope(results["gini"]) < 1.5,
        },
        "interpretation": (
            "With fixed k, the measured scaling is consistent with the "
            "O(n log n) WPM/Kolm implementations and the O(k n) Gini upper "
            "bound. Timing does not resolve the documented Gini optimality "
            "limitation."
        ),
    }
    summary["all_checks_passed"] = all(summary["checks"].values())
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()