"""Randomized quicksort, reservoir sampling, and one-sided amplification. Python 3.10+."""
from __future__ import annotations
from itertools import product
import random


def quicksort(values, rng):
    """Pre: finite mutually comparable values. Post: sorted permutation."""
    if len(values) < 2:
        return list(values)
    pivot = rng.choice(values)
    smaller = [value for value in values if value < pivot]
    equal = [value for value in values if value == pivot]
    larger = [value for value in values if value > pivot]
    # Recursive invariant: partitions preserve all values; bound: subproblem length.
    return quicksort(smaller, rng) + equal + quicksort(larger, rng)


def reservoir(stream, k, rng):
    """Pre: finite iterable, k>=0. Post: uniform subset of min(k,n) stream positions."""
    if k < 0:
        raise ValueError("k must be nonnegative")
    sample = []
    # Invariant: after i items, each position is included with probability min(1,k/i).
    # Bound: remaining items of the finite input.
    for i, value in enumerate(stream):
        if i < k:
            sample.append(value)
        else:
            chosen = rng.randrange(i + 1)
            if chosen < k:
                sample[chosen] = value
    return sample


def amplify_false_negatives(trials):
    """Pre: each True is sound, False may miss a yes. Post: accept if any trial accepts."""
    return any(trials)


def demo():
    rng = random.Random(7)
    print("Quicksort:", quicksort([7, 3, 9, 1, 5, 8, 2, 3], rng))
    print("Reservoir:", reservoir(range(20), 4, rng))
    print("False-negative amplification:", amplify_false_negatives([False, True, False]))
    print("19 independent runs reduce a 1/3 miss probability to", (1/3)**19)


def self_test():
    rng, count = random.Random(123), 0
    # Invariant: previous arrays sort correctly. Bound: remaining arrays of length <=5.
    for length in range(6):
        for values in product(range(3), repeat=length):
            check(quicksort(values, rng), sorted(values))
            count += 1
    for k in range(11):
        sample = reservoir(range(8), k, rng)
        check(len(sample), min(k, 8))
        check(len(set(sample)), len(sample))
        count += 2
    check(amplify_false_negatives([False, True, False]), True)
    check(amplify_false_negatives([False]*19), False)
    # Exact one-item uniformity: random index replacement rule gives every position 1/n.
    # Empirical frequencies are a diagnostic, not a proof.
    frequencies = [0]*5
    for _ in range(5000):
        frequencies[reservoir(range(5), 1, rng)[0]] += 1
    check(all(800 < value < 1200 for value in frequencies), True)
    return count + 3

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()