#!/usr/bin/env python3
"""Fresh exact CPU regimes protecting the existing Theorem 1/6 pages."""

from __future__ import annotations

import argparse
import json
import math
from fractions import Fraction
from pathlib import Path


def normalize(d: dict[int, Fraction]) -> dict[int, Fraction]:
    assert sum(d.values()) == 1
    return d


def policy_cost(d: dict[int, Fraction], k: int | None, b: int) -> Fraction:
    if k is None:
        return sum(Fraction(t) * p for t, p in d.items())
    return sum((Fraction(t) if t <= k else Fraction(k + b)) * p for t, p in d.items())


def optimum(d: dict[int, Fraction], b: int) -> tuple[int | None, Fraction]:
    candidates = [0, *sorted(d), None]
    values = [(policy_cost(d, k, b), k) for k in candidates]
    best = min(value for value, _ in values)
    tied = [k for value, k in values if value == best]
    return (None if None in tied else min(tied), best)


def tail(d: dict[int, Fraction], k: int) -> Fraction:
    return sum(p for t, p in d.items() if t > k)


def emd(first: dict[int, Fraction], second: dict[int, Fraction]) -> Fraction:
    points = sorted(set(first) | set(second))
    cumulative = Fraction(0)
    result = Fraction(0)
    for left, right in zip(points, points[1:]):
        cumulative += first.get(left, 0) - second.get(left, 0)
        result += abs(cumulative) * (right - left)
    return result


def protection_claim1() -> dict:
    # Non-power-of-two square buy costs, disjoint from the original 2^e grid.
    rows = []
    for root in (6, 10, 14, 18, 22, 26, 30, 34):
        b = root * root
        prediction = normalize({1: Fraction(1, 2), 2 * b: Fraction(1, 2)})
        khat, _ = optimum(prediction, b)
        u = next(k for k in [0, *sorted(prediction)] if tail(prediction, k) <= Fraction(1, root))
        kstar = min(khat + root if khat is not None else math.inf, u + root)
        truths = {
            "zero": prediction,
            "half": normalize({1: Fraction(1, 2), 2 * b + 1: Fraction(1, 2)}),
            "unit": normalize({1: Fraction(1, 2), 2 * b + 2: Fraction(1, 2)}),
            "root": normalize({1: Fraction(1, 2), 2 * b + 2 * root: Fraction(1, 2)}),
            "boundary": normalize({int(kstar) + 1: Fraction(1)}),
            "far": normalize({b**3: Fraction(1)}),
        }
        for name, truth in truths.items():
            eta = emd(prediction, truth)
            _, opt = optimum(truth, b)
            alg = policy_cost(truth, int(kstar), b)
            base = policy_cost(truth, int(khat + root), b)
            truncation_bound = Fraction(b, root) * (1 + eta)
            checks = []
            for a in sorted({0, *prediction, *truth, int(kstar), int(kstar + root)}):
                lhs = tail(truth, a + root)
                rhs = tail(prediction, a) + eta / root
                checks.append(lhs <= rhs)
            rows.append({
                "b": b,
                "family": name,
                "eta": str(eta),
                "kstar": int(kstar),
                "loss": str(alg - opt),
                "truncation_loss": str(alg - base),
                "truncation_bound": str(truncation_bound),
                "truncation_bound_holds": alg - base <= truncation_bound,
                "tail_checks": len(checks),
                "tail_checks_all_hold": all(checks),
            })
    return {
        "buy_costs": sorted({r["b"] for r in rows}),
        "cells": len(rows),
        "tail_checks": sum(r["tail_checks"] for r in rows),
        "all_tail_checks_hold": all(r["tail_checks_all_hold"] for r in rows),
        "all_truncation_bounds_hold": all(r["truncation_bound_holds"] for r in rows),
        "rows": rows,
    }


def two_phase(b: int) -> tuple[Fraction, Fraction, int, int, Fraction]:
    root = math.isqrt(b)
    assert root * root == b
    r1 = Fraction(b - 2, b)
    r2 = Fraction(2 * b - 1, 2 * b)
    transition = math.ceil(0.5 * b * math.log(b))

    def q(t: int) -> Fraction:
        return r1**t if t <= transition else r1**transition * r2 ** (t - transition)

    def cost(k: int) -> Fraction:
        if k <= transition:
            rent = sum(r1**i for i in range(k))
        else:
            rent = sum(r1**i for i in range(transition))
            rent += r1**transition * sum(r2**i for i in range(k - transition))
        return rent + b * q(k)

    u = 0
    while q(u) > Fraction(1, root):
        u += 1
    return r1, r2, transition, u, cost


def protection_claims2_and4() -> dict:
    rows = []
    # Fresh square scales, separate from both prior power-of-two runs.
    for root in (6, 10, 14, 18, 22, 26, 30, 34):
        b = root * root
        r1, r2, transition, u, cost = two_phase(b)
        chosen = u + root
        identity_errors = []
        for k in (0, 1, max(0, u - 1), u, transition, chosen):
            lhs = cost(k + 1) - cost(k)
            rhs = (r1**k if k <= transition else r1**transition * r2 ** (k - transition)) - b * (
                (r1**k if k <= transition else r1**transition * r2 ** (k - transition))
                - (r1 ** (k + 1) if k + 1 <= transition else r1**transition * r2 ** (k + 1 - transition))
            )
            identity_errors.append(lhs - rhs)
        consistency = cost(chosen) - cost(transition)
        # Theorem 6's exact threshold excess, with C=1.
        target = Fraction(2) * root
        lo, hi = 0, transition
        while lo < hi:
            mid = (lo + hi + 1) // 2
            excess = Fraction(b, 2) * (r1**mid - r1**transition)
            if excess >= target:
                lo = mid
            else:
                hi = mid - 1
        rows.append({
            "b": b,
            "transition": transition,
            "chosen_threshold": chosen,
            "consistency_over_sqrt_b": float(consistency / root),
            "robustness_over_b_log_b": float(Fraction(chosen, 1) / (b * math.log(b))),
            "theorem6_B0": lo,
            "theorem6_randomized_ratio": float(Fraction(lo, 2) / (b * math.log(b))),
            "max_identity_error": str(max(abs(error) for error in identity_errors)),
        })
    return {
        "scales": [r["b"] for r in rows],
        "rows": rows,
        "all_policy_cost_identities_exact": all(r["max_identity_error"] == "0" for r in rows),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = {
        "claim1_fresh_exact": protection_claim1(),
        "claim2_and_claim4_fresh_exact": protection_claims2_and4(),
    }
    args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()