#!/usr/bin/env python3
"""CPU-only scope expansion for the generalized-convex audit.

This file deliberately performs new work rather than copying the original
eight-point result.  The finite-transform, leanness, auction, and transport
checks are all evaluated again on larger or denser exact grids.  The script
writes only the requested JSON report; it does not modify the packaged
artifacts used by the original reproduction.
"""

from __future__ import annotations

import argparse
import itertools
import json
import re
import sys
from fractions import Fraction
from pathlib import Path


def finite_transform(x: Fraction, ys: list[Fraction]) -> Fraction:
    return max(x * y - y * y / 2 for y in ys)


def transform_scope() -> dict:
    rows = []
    for k in (128, 256, 512):
        h = Fraction(1, k)
        ys = [Fraction(j, k) for j in range(-k, k + 1)]
        errors = []
        gradient_errors = []
        gradient_checks = 0
        for left, right in zip(ys, ys[1:]):
            midpoint = (left + right) / 2
            x = midpoint
            errors.append(x * x / 2 - finite_transform(x, ys))
            for offset in (Fraction(-1, 1000), Fraction(-1, 4), Fraction(1, 4), Fraction(1, 1000)):
                sample = midpoint + offset * h
                values = [sample * y - y * y / 2 for y in ys]
                best = max(range(len(ys)), key=values.__getitem__)
                assert values.count(values[best]) == 1
                gradient_errors.append(abs(ys[best] - sample))
                gradient_checks += 1
        maximum = max(errors)
        observed_gradient = max(gradient_errors)
        assert min(errors) >= 0 and maximum == h * h / 8
        assert observed_gradient == Fraction(499, 1000) * h
        rows.append(
            {
                "k": k,
                "finite_atoms": len(ys),
                "cells_exhausted": len(errors),
                "exact_uniform_error": f"{maximum.numerator}/{maximum.denominator}",
                "uniform_error": float(maximum),
                "gradient_points_checked": gradient_checks,
                "exact_gradient_error": f"{observed_gradient.numerator}/{observed_gradient.denominator}",
                "gradient_error": float(observed_gradient),
            }
        )
    return {
        "new_k_values": [128, 256, 512],
        "rows": rows,
        "claim1_scope": "all cells of three grids through 1,025 finite atoms",
        "claim2_scope": "four non-tie points in every cell of the same three grids",
    }


def lean_scope() -> dict:
    ys = [Fraction(j, 4) for j in range(-8, 9)]
    a_values = [Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(3, 2), Fraction(2), Fraction(4)]
    b_values = [Fraction(-2), Fraction(-1), Fraction(-1, 2), Fraction(0), Fraction(1, 2), Fraction(1), Fraction(2)]
    parameters = [(a, b, (a + b) / 17) for a in a_values for b in b_values]
    pair_count = 0
    combination_count = 0
    inequality_count = 0
    for p, q in itertools.combinations(parameters, 2):
        pair_count += 1
        for lam in (Fraction(1, 5), Fraction(2, 5), Fraction(3, 5), Fraction(4, 5)):
            a = (1 - lam) * p[0] + lam * q[0]
            b = (1 - lam) * p[1] + lam * q[1]
            c = (1 - lam) * p[2] + lam * q[2]
            assert a > 0
            for yi in ys:
                x = 2 * a * yi + b
                chosen = x * yi - (a * yi * yi + b * yi + c)
                for yj in ys:
                    competitor = x * yj - (a * yj * yj + b * yj + c)
                    assert chosen - competitor >= 0
                    inequality_count += 1
            combination_count += 1
    return {
        "a_values": [str(x) for x in a_values],
        "b_values": [str(x) for x in b_values],
        "y_grid": [str(x) for x in ys],
        "parameterizations": len(parameters),
        "unordered_pairs": pair_count,
        "convex_combinations_checked": combination_count,
        "exact_activation_inequalities": inequality_count,
    }


def transport_scope() -> dict:
    rows = []
    for n in (8, 9, 10):
        denominator = n * (n - 1) ** 2
        best_numerator = -1
        best_perm = None
        permutations_checked = 0
        for perm in itertools.permutations(range(n)):
            numerator = sum(i * perm[i] for i in range(n))
            if numerator > best_numerator:
                best_numerator = numerator
                best_perm = perm
            permutations_checked += 1
        identity = tuple(range(n))
        identity_numerator = sum(i * i for i in range(n))
        reverse = tuple(reversed(identity))
        reverse_numerator = sum(i * reverse[i] for i in range(n))
        dual_slacks = [(i - j) ** 2 for i in range(n) for j in range(n)]
        assert best_perm == identity
        assert best_numerator == identity_numerator
        assert min(dual_slacks) == 0
        rows.append(
            {
                "grid_points": n,
                "permutations_exhausted": permutations_checked,
                "optimal_permutation": list(best_perm),
                "primal_numerator": best_numerator,
                "dual_numerator": identity_numerator,
                "exact_primal_dual_gap": "0/1",
                "dual_constraints_checked": len(dual_slacks),
                "minimum_dual_slack_numerator": min(dual_slacks),
                "reverse_numerator": reverse_numerator,
                "reverse_gap_numerator": identity_numerator - reverse_numerator,
                "normalization_denominator": denominator,
            }
        )
    return {
        "surplus": "Phi(x,y)=xy",
        "grid_sizes": [8, 9, 10],
        "new_larger_grids": [9, 10],
        "rows": rows,
        "interpretation": "The identity map remains exactly optimal after exhaustive enumeration through 10 points; dual equality and zero minimum slack hold on every grid.",
    }


def auction_scope(author_root: Path) -> dict:
    sys.path.insert(0, str(author_root))
    import torch
    from mech_design.mechanism import Mechanism

    def dot_kernel(x, y):
        return (x * y).sum(dim=-1)

    model = Mechanism(
        npoints=1,
        kernel=dot_kernel,
        y_dim=1,
        temp=1.0,
        is_Y_parameter=False,
        is_there_default=True,
        y_min=0.0,
        y_max=1.0,
    )
    with torch.no_grad():
        model.Y_rest_raw.fill_(1.0)
        model.intercept_rest.fill_(0.5)
    rows = []
    with torch.no_grad():
        for points in (16001, 32001):
            xs = torch.linspace(0.0, 1.0, points)[:, None]
            choices, _ = model.forward(xs, selection_mode="hard")
            allocation = choices[:, 0]
            below = allocation[xs[:, 0] < 0.499]
            above = allocation[xs[:, 0] > 0.501]
            assert float(below.max()) == 0.0 and float(above.min()) == 1.0
            rows.append(
                {
                    "grid_points": points,
                    "max_allocation_below_0.499": float(below.max()),
                    "min_allocation_above_0.501": float(above.min()),
                    "below_points": int(below.numel()),
                    "above_points": int(above.numel()),
                }
            )
    return {
        "author_commit": "85a5da444a146ea28945173e22e3d130163dae42",
        "fresh_grid_rows": rows,
        "interpretation": "The released Mechanism preserves the exact threshold on two additional dense CPU grids, independently of the original 4,001-point run.",
    }


def table_scope(source_root: Path) -> dict:
    """Re-read every registered table row and recompute the printed gaps."""
    text = (source_root / "src/sections/_VII_experiments.tex").read_text(encoding="utf-8")
    pattern = re.compile(r"^(1|2|5|10|20)\s*&\s*([0-9.]+)\s*&\s*(---|[0-9.]+)", re.MULTILINE)
    rows = []
    for n, learned, benchmark in pattern.findall(text):
        gap = None if benchmark == "---" else abs(Fraction(learned) - Fraction(benchmark))
        rows.append(
            {
                "n": int(n),
                "learned": learned,
                "straight_jacket": benchmark,
                "exact_gap": None if gap is None else f"{gap.numerator}/{gap.denominator}",
                "within_printed_0.001": None if gap is None else gap <= Fraction(1, 1000),
            }
        )
    assert [row["n"] for row in rows] == [1, 2, 5, 10, 20]
    assert [row["within_printed_0.001"] for row in rows if row["within_printed_0.001"] is not None] == [True, True, True, True]
    return {
        "rows_reparsed": rows,
        "comparable_rows": 4,
        "exact_match_rows": [row["n"] for row in rows if row["exact_gap"] == "0/1"],
        "nonzero_exact_gap_rows": [row["n"] for row in rows if row["exact_gap"] not in (None, "0/1")],
        "scope": "all five registered n rows, with exact Fraction arithmetic and a separate 0.001 printed-precision check",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--author-root", type=Path, required=True)
    parser.add_argument("--source-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = {
        "transform": transform_scope(),
        "lean": lean_scope(),
        "transport": transport_scope(),
        "auction": auction_scope(args.author_root),
        "table": table_scope(args.source_root),
    }
    args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"transform_rows": len(report["transform"]["rows"]), "transport_rows": len(report["transport"]["rows"]), "auction_rows": len(report["auction"]["fresh_grid_rows"])}, sort_keys=True))


if __name__ == "__main__":
    main()