#!/usr/bin/env python3
"""Exact finite audit of Theorem 4 (the matching EMD lower bound).

The source construction has two truths.  They differ only by moving mass
epsilon from day 4 to day 2, while a midpoint prediction is equally close to
both.  A deterministic ski-rental policy is a threshold K (or infinity).
Enumerating every threshold proves the two incompatible threshold regions;
the final mixture inequality is Yao's argument for randomized policies.
"""

from __future__ import annotations

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


INF = None


def distribution_pair(b: int) -> tuple[dict[int, Fraction], dict[int, Fraction], dict[int, Fraction], Fraction]:
    delta = Fraction(1, 8 * b)
    epsilon = Fraction(1, 2 * b) + delta
    p1 = {
        1: Fraction(1, 2),
        2: Fraction(1, 4),
        4: epsilon,
        b + 3: Fraction(1, 4) - epsilon,
    }
    p2 = {
        1: Fraction(1, 2),
        2: Fraction(1, 4) + epsilon,
        b + 3: Fraction(1, 4) - epsilon,
    }
    midpoint = {
        1: Fraction(1, 2),
        2: Fraction(1, 4) + epsilon / 2,
        4: epsilon / 2,
        b + 3: Fraction(1, 4) - epsilon,
    }
    return p1, p2, midpoint, epsilon


def policy_cost(distribution: dict[int, Fraction], threshold: int | None, b: int) -> Fraction:
    if threshold is INF:
        return sum(Fraction(day) * mass for day, mass in distribution.items())
    return sum(
        (Fraction(day) if day <= threshold else Fraction(threshold + b)) * mass
        for day, mass in distribution.items()
    )


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


def audit_one(b: int) -> dict:
    p1, p2, midpoint, epsilon = distribution_pair(b)
    candidates: list[int | None] = list(range(0, b + 4)) + [INF]
    costs1 = {k: policy_cost(p1, k, b) for k in candidates}
    costs2 = {k: policy_cost(p2, k, b) for k in candidates}
    opt1 = min(costs1.values())
    opt2 = min(costs2.values())
    opt1_thresholds = [k for k, value in costs1.items() if value == opt1]
    opt2_thresholds = [k for k, value in costs2.items() if value == opt2]
    # Buying after the last support point ties rent-forever; the source's
    # canonical choice is A_infinity and it is present in the optimum set.
    assert INF in opt1_thresholds, opt1_thresholds
    assert opt2_thresholds == [2], opt2_thresholds

    low = [costs1[k] - opt1 for k in range(4)]
    high = [costs2[k] - opt2 for k in candidates if k is INF or k >= 4]
    min_low = min(low)
    min_high = min(high)
    randomized_worst_case = min(min_low, min_high) / 2
    midpoint_emd_1 = emd(midpoint, p1)
    midpoint_emd_2 = emd(midpoint, p2)
    assert midpoint_emd_1 == epsilon == midpoint_emd_2
    # For q=P(threshold <= 3), p1 loss >= q*min_low and p2 loss >=
    # (1-q)*min_high.  Their maximum is at least half the smaller constant.
    assert randomized_worst_case > 0
    return {
        "b": b,
        "delta": str(Fraction(1, 8 * b)),
        "epsilon": str(epsilon),
        "support_max": b + 3,
        "prediction_emd_to_p1": str(midpoint_emd_1),
        "prediction_emd_to_p2": str(midpoint_emd_2),
        "p1_optimal_threshold": "infinity",
        "p2_optimal_threshold": 2,
        "low_thresholds_0_to_3_loss": [str(value) for value in low],
        "min_loss_p1_for_K_le_3": str(min_low),
        "min_loss_p2_for_K_ge_4_or_infinity": str(min_high),
        "randomized_worst_case_lower_bound": str(randomized_worst_case),
        "randomized_bound_over_b_times_emd": float(randomized_worst_case / (b * epsilon)),
        "enumerated_deterministic_policies": len(candidates),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    # These are finite support instances, with a nontrivial range of b.
    rows = [audit_one(b) for b in (16, 32, 64, 128, 256, 1024, 4096, 16384, 65536, 262144)]
    result = {
        "theorem": "Theorem 4",
        "arithmetic": "exact Fraction",
        "rows": rows,
        "all_midpoint_emd_equal_epsilon": all(
            r["prediction_emd_to_p1"] == r["prediction_emd_to_p2"] == r["epsilon"] for r in rows
        ),
        "all_optima_match_source": all(
            r["p1_optimal_threshold"] == "infinity" and r["p2_optimal_threshold"] == 2 for r in rows
        ),
        "all_randomized_lower_bounds_positive": all(
            Fraction(r["randomized_worst_case_lower_bound"]) > 0 for r in rows
        ),
    }
    args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()