#!/usr/bin/env python3
"""Fresh CPU-only protection runs for the three thin, already-good claims."""
from __future__ import annotations

import hashlib
import json
import math
import re
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
ALGORITHM_LABELS = (
    "alg:sampler",
    "alg:GF-DRO",
    "alg:SDRO-NGD",
    "alg:SDRO-WFR",
    "alg:SDRO-SVG",
    "alg:SDRO_rgo",
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def algorithm_inventory() -> dict:
    rows = []
    for version in ("v1", "current"):
        text = (ROOT / f"source_{version}" / "main.tex").read_text(encoding="utf-8")
        counts = {label: len(re.findall(r"label\{" + re.escape(label) + r"\}", text)) for label in ALGORITHM_LABELS}
        rows.append({"version": version, "label_counts": counts, "six_present": all(counts.values())})
    return {
        "versions": rows,
        "all_12_label_occurrences_present": all(row["six_present"] for row in rows),
        "v1_source_sha256": sha256(ROOT / "source_v1.tar.gz"),
        "current_source_sha256": sha256(ROOT / "source_current.tar.gz"),
    }


def flow_time_grid() -> dict:
    rows = []
    for lam in (1 / 16, 1 / 3, 3 / 4, 3 / 2, 3):
        for lipschitz in (1 / 3, 2 / 3, 3, 5):
            for epsilon in (1 / 32, 1 / 128, 1 / 512, 1 / 2048):
                initial = lipschitz / math.sqrt(lam)
                threshold = math.log(initial / epsilon) / lam
                at_threshold = initial * math.exp(-lam * threshold)
                early = initial * math.exp(-lam * 0.8 * threshold)
                rows.append(
                    {
                        "lambda": lam,
                        "L": lipschitz,
                        "epsilon": epsilon,
                        "threshold_time": threshold,
                        "error_at_threshold": at_threshold,
                        "error_at_80pct": early,
                        "threshold_ratio_error": at_threshold / epsilon - 1,
                        "early_control_fails": early > epsilon,
                    }
                )
    return {
        "cells": len(rows),
        "max_abs_threshold_ratio_error": max(abs(row["threshold_ratio_error"]) for row in rows),
        "all_thresholds_pass": all(abs(row["threshold_ratio_error"]) < 1e-12 for row in rows),
        "all_early_controls_fail": all(row["early_control_fails"] for row in rows),
        "sample_rows": [rows[0], rows[len(rows) // 2], rows[-1]],
    }


def complexity_grid() -> dict:
    rows = []
    for epsilon in (1 / 2, 1 / 5, 1 / 13, 1 / 37, 1 / 101):
        for dimension in (3, 7, 19, 53):
            for l_phi, l_u, l_f, lam_u in ((1 / 3, 2 / 3, 5 / 4, 1 / 2), (3, 2, 1 / 5, 3 / 2)):
                prefactor = l_phi * l_u**2 * l_f**2 * dimension**2 / lam_u**3
                total = prefactor * epsilon**-4 * math.log(1 / epsilon)
                rows.append(
                    {
                        "epsilon_opt": epsilon,
                        "dimension": dimension,
                        "L_Phi": l_phi,
                        "L_U": l_u,
                        "L_f": l_f,
                        "lambda_U": lam_u,
                        "total": total,
                        "normalized": total / (prefactor * epsilon**-4 * math.log(1 / epsilon)),
                    }
                )
    eps = sorted({row["epsilon_opt"] for row in rows})
    slope = sum((math.log(x) - sum(math.log(y) for y in eps) / len(eps)) * (math.log(x**-4) - sum(math.log(y**-4) for y in eps) / len(eps)) for x in eps) / sum((math.log(x) - sum(math.log(y) for y in eps) / len(eps)) ** 2 for x in eps)
    return {
        "cells": len(rows),
        "max_normalized_identity_error": max(abs(row["normalized"] - 1) for row in rows),
        "epsilon_polynomial_slope": slope,
        "slope_error_from_minus_4": slope + 4,
        "sample_rows": [rows[0], rows[len(rows) // 2], rows[-1]],
    }


def main() -> None:
    result = {
        "algorithm_inventory": algorithm_inventory(),
        "flow_time_grid": flow_time_grid(),
        "complexity_grid": complexity_grid(),
    }
    result["all_protection_checks_pass"] = (
        result["algorithm_inventory"]["all_12_label_occurrences_present"]
        and result["flow_time_grid"]["cells"] == 80
        and result["flow_time_grid"]["all_thresholds_pass"]
        and result["flow_time_grid"]["all_early_controls_fail"]
        and result["complexity_grid"]["cells"] == 40
        and result["complexity_grid"]["max_normalized_identity_error"] < 1e-12
        and abs(result["complexity_grid"]["slope_error_from_minus_4"]) < 1e-12
    )
    print(json.dumps(result, indent=2, sort_keys=True))
    if not result["all_protection_checks_pass"]:
        raise SystemExit("fresh protection check failed")


if __name__ == "__main__":
    main()