#!/usr/bin/env python3
"""Exact arithmetic for two explanatory articles; no model or benchmark is run."""

from decimal import Decimal
from fractions import Fraction
import json


def corrected_distribution(target, draft):
    """Enumerate the probability mass of one speculative-sampling position."""
    assert set(target) == set(draft)
    assert sum(target.values()) == sum(draft.values()) == 1
    assert all(x >= 0 for x in (*target.values(), *draft.values()))
    accepted = {token: min(target[token], draft[token]) for token in target}
    rejection = 1 - sum(accepted.values())
    residual_mass = {token: max(target[token] - draft[token], 0) for token in target}
    assert sum(residual_mass.values()) == rejection
    residual = (
        {token: residual_mass[token] / rejection for token in target}
        if rejection else {token: Fraction(0) for token in target}
    )
    output = {token: accepted[token] + rejection * residual[token] for token in target}
    return accepted, rejection, residual, output


def main():
    # Hypothetical full K/V representation; not an MLA model or observed traffic.
    kv_bytes = 2 * 32 * 8192 * 8 * 128 * 2
    assert kv_bytes == 1024**3
    transfer_ms = Decimal(kv_bytes * 8) / Decimal(400_000_000_000) * 1000
    assert transfer_ms == Decimal('21.47483648')

    p = {'A': Fraction(3, 5), 'B': Fraction(2, 5)}
    q = {'A': Fraction(4, 5), 'B': Fraction(1, 5)}
    accepted, rejection, residual, output = corrected_distribution(p, q)
    assert accepted == {'A': Fraction(3, 5), 'B': Fraction(1, 5)}
    assert rejection == Fraction(1, 5)
    assert residual == {'A': Fraction(0), 'B': Fraction(1)}
    assert output == p

    # Include equal and disjoint distributions, as well as zero-mass symbols.
    checked = 0
    for p_a in range(11):
        for q_a in range(11):
            target = {'A': Fraction(p_a, 10), 'B': 1 - Fraction(p_a, 10)}
            draft = {'A': Fraction(q_a, 10), 'B': 1 - Fraction(q_a, 10)}
            assert corrected_distribution(target, draft)[3] == target
            checked += 1
    assert checked == 121

    # A greedy illustration only; this does not implement a Transformer.
    proposals = ['A', 'B', 'C', 'D']
    accepted_prefix = proposals[:2]
    emitted = accepted_prefix + ['X']
    assert emitted == ['A', 'B', 'X']
    # X is sampled; its own target K/V may be materialized by the next pass.

    # Dated, isolated Sonnet 5.5 list-price arithmetic, USD per million tokens.
    generated = Decimal(1000) * Decimal(10) / 1_000_000
    next_write = Decimal(1000) * Decimal('2.50') / 1_000_000
    next_read = Decimal(1000) * Decimal('0.10') / 1_000_000
    assert generated + next_write == Decimal('0.0125')
    assert next_read == Decimal('0.0001')
    print(json.dumps({
        'kind': 'synthetic exact-arithmetic checks, not an inference benchmark',
        'model_calls': 0,
        'hypothetical_kv_bytes': kv_bytes,
        'raw_400_gbit_serialization_floor_ms': str(transfer_ms),
        'synthetic_output_probability': {k: str(v) for k, v in output.items()},
        'two_symbol_distribution_pairs_checked': checked,
        'greedy_example_emitted_tokens': emitted,
        'isolated_generated_then_written_usd': str(generated + next_write),
        'counterfactual_next_read_usd': str(next_read),
        'pricing_checked_on': '2026-11-11',
    }, indent=2))


if __name__ == '__main__':
    main()