# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib==3.10.8"]
# ///
"""Check the exact samplers and reproduce coin_tree.png and bit_cost.png.

Run with `uv run --locked gen_figures.py`. Download figure_quality.py into
the parent directory when running the public reproduction bundle standalone.
"""

from __future__ import annotations

import math
import sys
from collections.abc import Callable
from fractions import Fraction
from pathlib import Path

import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure


def bernoulli(a: int, b: int, flip: Callable[[], int]) -> int:
    """Sample Bernoulli(a/b); flip must supply independent fair 0/1 bits."""
    if not isinstance(a, int) or not isinstance(b, int):
        raise TypeError("a and b must be integers")
    if b <= 0 or not 0 <= a <= b:
        raise ValueError("require b > 0 and 0 <= a <= b")
    if a == b:
        return 1
    while a:
        digit, a = divmod(2 * a, b)
        bit = flip()
        if bit not in (0, 1):
            raise ValueError("flip must return 0 or 1")
        if bit == 1:
            return digit
    return 0


def categorical(weights: tuple[int, ...], flip: Callable[[], int]) -> int:
    """Sample a category using exact nonnegative integer weights."""
    if not weights or any(not isinstance(w, int) or w < 0 for w in weights):
        raise ValueError("require nonnegative integer weights")
    remaining = sum(weights)
    if remaining == 0:
        raise ValueError("at least one weight must be positive")
    for index, weight in enumerate(weights):
        if bernoulli(weight, remaining, flip):
            return index
        remaining -= weight
    raise AssertionError("the final positive category must be selected")


def uniform_below(n: int, flip: Callable[[], int]) -> int:
    """Fast Dice Roller: retain the uniform remainder after rejection."""
    if n < 1:
        raise ValueError("n must be positive")
    if n == 1:
        return 0
    v, c = 1, 0
    while True:
        v, c = 2 * v, 2 * c + flip()
        if v >= n:
            if c < n:
                return c
            v, c = v - n, c - n


def inverse_e(flip: Callable[[], int]) -> int:
    """An exact 1/e coin built from rational coins, without digits of e."""
    n = 2
    while bernoulli(1, n, flip):
        n += 1
    return int((n - 1) % 2 == 0)


class NeedBit(Exception):
    """A deterministic replay reached the end of its supplied bit prefix."""


def prefix_tree(sample: Callable, depth: int) -> tuple[dict, dict, list]:
    """Explore the actual sampler, recording exact leaf and unresolved mass."""
    active = [()]
    masses: dict[int, Fraction] = {}
    leaves: dict[int, Fraction] = {}
    tails = []
    for level in range(depth + 1):
        unresolved = []
        for prefix in active:
            bits = iter(prefix)
            consumed = 0

            def flip(bits=bits) -> int:
                nonlocal consumed
                try:
                    bit = next(bits)
                except StopIteration:
                    raise NeedBit from None
                consumed += 1
                return bit

            try:
                outcome = sample(flip)
            except NeedBit:
                unresolved.append(prefix)
            else:
                assert consumed == level
                mass = Fraction(1, 2**level)
                masses[outcome] = masses.get(outcome, Fraction()) + mass
                leaves[level] = leaves.get(level, Fraction()) + mass
        tail = Fraction(len(unresolved), 2**level)
        assert sum(masses.values()) + tail == 1
        tails.append(tail)
        active = [prefix + (bit,) for prefix in unresolved for bit in (0, 1)]
    return masses, leaves, tails


def optimal_cost(p: Fraction) -> Fraction:
    if p in (0, 1):
        return Fraction()
    denominator = p.denominator
    if denominator & (denominator - 1):
        return Fraction(2)
    return 2 - Fraction(2, denominator)


def verify() -> None:
    depth = 12
    cases = 0
    for denominator in range(1, 65):
        for numerator in range(denominator + 1):
            p = Fraction(numerator, denominator)
            masses, leaves, tails = prefix_tree(
                lambda flip, a=numerator, b=denominator: bernoulli(a, b, flip), depth
            )
            success = masses.get(1, Fraction())
            unresolved = tails[-1]
            assert success <= p <= success + unresolved
            # At depth k, no exact tree can have assigned more than
            # floor(2^k p) success cells or floor(2^k (1-p)) failure cells.
            for k, tail in enumerate(tails):
                cells = 2**k
                bound = 1 - Fraction(
                    (cells * p).__floor__() + (cells * (1 - p)).__floor__(),
                    cells,
                )
                assert tail == bound, (p, k, tail, bound)
            if not unresolved:
                assert success == p
                assert sum(k * mass for k, mass in leaves.items()) == optimal_cost(p)
            else:
                assert sum(tails[:-1]) == 2 - Fraction(2, 2**depth)
            cases += 1
    for weights in ((1, 2, 3), (0, 1, 0), (3, 1, 4, 2), (1,), (1, 1, 1, 1)):
        masses, _, tails = prefix_tree(
            lambda flip, weights=weights: categorical(weights, flip), depth
        )
        for i, weight in enumerate(weights):
            exact = Fraction(weight, sum(weights))
            lower = masses.get(i, Fraction())
            assert lower <= exact <= lower + tails[-1]
    for a, b in ((-1, 2), (3, 2), (0, 0), (1, -1)):
        try:
            bernoulli(a, b, lambda: 0)
        except ValueError:
            pass
        else:
            raise AssertionError("invalid probability accepted")
    for n in (1, 2, 3, 5, 6, 7, 16):
        masses, _, tails = prefix_tree(lambda flip, n=n: uniform_below(n, flip), depth)
        assert len(set(masses.values())) == 1
        for k, tail in enumerate(tails):
            assert tail == Fraction(2**k % n, 2**k)
        assert all(
            mass <= Fraction(1, n) <= mass + tails[-1] for mass in masses.values()
        )
    _, _, four_tails = prefix_tree(lambda flip: categorical((1, 1, 1, 1), flip), depth)
    assert abs(sum(four_tails[:-1]) - Fraction(7, 2)) < Fraction(1, 100)
    assert Fraction(1, 4) / (1 - Fraction(1, 4)) == Fraction(1, 3)
    assert optimal_cost(Fraction(3, 8)) == Fraction(7, 4)
    e_masses, _, e_tails = prefix_tree(inverse_e, 18)
    lower_e = sum(Fraction((-1) ** j, math.factorial(j)) for j in range(22))
    upper_e = lower_e + Fraction(1, math.factorial(21))
    assert e_masses[1] <= lower_e < upper_e <= e_masses[1] + e_tails[-1]
    e_cost = sum(
        optimal_cost(Fraction(1, n)) / math.factorial(n - 1) for n in range(2, 22)
    )
    print(
        f"1/e coin: success enclosed by [{float(e_masses[1]):.8f}, {float(e_masses[1] + e_tails[-1]):.8f}]"
    )
    print(f"1/e coin expected fair flips: {float(e_cost):.8f} (series through n=21)")
    print(
        f"PASS: {cases} rational inputs; exact decision-tree tails through depth {depth}"
    )
    print("PASS: categorical probability enclosures, endpoints, invalid inputs")
    print("PASS: recycled dice have equal leaf masses and optimal uniform tail bounds")
    print("p         optimal expected flips")
    for p in (
        Fraction(),
        Fraction(1, 8),
        Fraction(1, 3),
        Fraction(3, 8),
        Fraction(1, 2),
        Fraction(1),
    ):
        print(f"{p!s:8}  {optimal_cost(p)}")


def figures() -> None:
    configure_matplotlib()
    plt.rcParams.update(
        {
            "font.size": 14,
            "axes.labelsize": 14,
            "xtick.labelsize": 13,
            "ytick.labelsize": 13,
        }
    )
    ink, blue, rust = "#292929", "#235c80", "#934428"

    fig, ax = plt.subplots(figsize=(4.4, 4.0), layout="constrained")
    ax.set_xlim(-0.35, 3.8)
    ax.set_ylim(-0.25, 3.25)
    ax.axis("off")
    for k in range(3):
        y = 2.7 - k
        ax.text(
            0,
            y,
            f"Flip {k + 1}",
            ha="center",
            va="center",
            color=ink,
            bbox={"boxstyle": "round,pad=0.4", "fc": "white", "ec": ink},
        )
        ax.annotate(
            "",
            xy=(2.25, y),
            xytext=(0.55, y),
            arrowprops={"arrowstyle": "->", "color": blue},
        )
        ax.text(1.35, y + 0.13, "H", ha="center", color=blue)
        ax.text(
            2.4,
            y,
            f"return {k % 2}\nmass 1/{2 ** (k + 1)}",
            va="center",
            color=ink,
        )
        ax.annotate(
            "",
            xy=(0, y - 0.7),
            xytext=(0, y - 0.2),
            arrowprops={"arrowstyle": "->", "color": rust},
        )
        ax.text(0.15, y - 0.47, "T", color=rust)
    ax.text(0, -0.18, "…", ha="center", color=ink)
    validate_figure(fig)
    fig.savefig(figure_output_path(__file__, "coin_tree.png"), dpi=300)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(4.4, 4.4), layout="constrained")
    xs = [i / 500 for i in range(1, 500)]
    entropy = [-p * math.log2(p) - (1 - p) * math.log2(1 - p) for p in xs]
    ax.plot(xs, entropy, color=blue, label="Entropy H(p)")
    ax.axhline(2, color=ink, linestyle="--", label="Nondyadic: 2 flips")
    for m in range(1, 6):
        ps = [Fraction(a, 2**m) for a in range(1, 2**m, 2)]
        ax.scatter(
            [float(p) for p in ps],
            [float(optimal_cost(p)) for p in ps],
            color=rust,
            s=28,
            label="Dyadic probabilities" if m == 1 else None,
            zorder=3,
        )
    ax.scatter([0, 1], [0, 0], color=rust, s=28, zorder=3)
    ax.set(
        xlabel="Probability p",
        ylabel="Bits per output",
        xlim=(-0.02, 1.02),
        ylim=(-0.08, 2.18),
    )
    ax.legend(loc="lower center", fontsize=13, frameon=False)
    ax.spines[["top", "right"]].set_visible(False)
    validate_figure(fig)
    fig.savefig(figure_output_path(__file__, "bit_cost.png"), dpi=300)
    plt.close(fig)


if __name__ == "__main__":
    verify()
    figures()