"""Reversible RLE, exact arithmetic intervals, and real stdlib codecs. Python 3.10+."""
from __future__ import annotations
from fractions import Fraction
from itertools import groupby
import bz2
import lzma
import zlib


def rle_encode(data: bytes):
    """Pre: bytes. Post: runs (positive count, byte), preserving all input values."""
    return [(sum(1 for _ in group), value) for value, group in groupby(data)]


def rle_decode(runs):
    """Pre: positive integer counts and byte values. Post: exactly reconstructed bytes."""
    runs = list(runs)
    if any(not isinstance(n, int) or n <= 0 or not isinstance(value, int)
           or not 0 <= value <= 255 for n, value in runs):
        raise ValueError("runs require positive integer counts and bytes")
    return b"".join(bytes([value])*n for n, value in runs)


MODEL = {"A": (Fraction(0), Fraction(1, 2)),
         "B": (Fraction(1, 2), Fraction(4, 5)),
         "C": (Fraction(4, 5), Fraction(1))}


def arithmetic_interval(text):
    """Pre: only A/B/C. Post: exact subinterval representing this text under MODEL."""
    low, high = Fraction(0), Fraction(1)
    # Invariant: interval represents the complete consumed prefix. Bound: remaining symbols.
    for symbol in text:
        a, b = MODEL[symbol]
        width = high-low
        low, high = low+width*a, low+width*b
    return low, high


def arithmetic_decode(value, length):
    """Pre: 0<=value<1, length>=0. Post: exactly 'length' symbols under the same model."""
    if not 0 <= value < 1 or length < 0:
        raise ValueError("value in [0,1) and nonnegative length required")
    low, high, output = Fraction(0), Fraction(1), []
    # Invariant: output is the decoded prefix; current interval contains value.
    # Bound: length-len(output).
    for _ in range(length):
        position = (value-low)/(high-low)
        symbol = next(s for s, (a, b) in MODEL.items() if a <= position < b)
        a, b = MODEL[symbol]
        width = high-low
        low, high = low+width*a, low+width*b
        output.append(symbol)
    return "".join(output)


def demo():
    data = b"AAAAABBBCCAAAA"
    print("RLE runs:", rle_encode(data))
    interval = arithmetic_interval("BAC")
    print("BAC interval:", interval)
    print("Decode midpoint:", arithmetic_decode(sum(interval)/2, 3))
    corpus = b"the cat sat on the mat. the dog sat on the mat.\n"*200
    for name, compress, decompress in (("zlib", zlib.compress, zlib.decompress),
                                      ("bz2", bz2.compress, bz2.decompress),
                                      ("lzma", lzma.compress, lzma.decompress)):
        encoded = compress(corpus)
        print(name, "raw bytes:", len(corpus), "encoded bytes:", len(encoded),
              "round trip:", decompress(encoded) == corpus)
    print("Exact Fractions explain intervals; production coders use finite-word renormalization.")


def self_test():
    count = 0
    # Invariant: preceding payloads round-trip. Bound: remaining fixtures/codecs.
    for data in (b"", b"a", b"a"*1000, bytes(range(256)), b"abracadabra"*20):
        check(rle_decode(rle_encode(data)), data)
        count += 1
        for compress, decompress in ((zlib.compress, zlib.decompress),
                                     (bz2.compress, bz2.decompress), (lzma.compress, lzma.decompress)):
            check(decompress(compress(data)), data)
            count += 1
    for text in ("", "A", "B", "C", "BAC", "ABCABC", "AAAA"):
        low, high = arithmetic_interval(text)
        check(arithmetic_decode((low+high)/2, len(text)), text)
        count += 1
    check(arithmetic_interval("BAC"), (Fraction(31, 50), Fraction(13, 20)))
    check(rle_decode(iter([(2, 65)])), b"AA")
    return count + 2

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()