# Contains all of the code changes locally, so can be run on unmodified main branch or 3.10 Python

from random import shuffle, Random, BPF, _floor
from time import time


RANDOM_TRIAL_COUNT = 4000
INCLUDE_WITHOUT_GETRANDBITS = False


class Main(Random):
    # Current main _randbelow_with_getrandbits()
    def _randbelow(self, n):
        if not n:
            return 0
        getrandbits = self.getrandbits
        k = n.bit_length()  # don't use (n-1) here because n can be 1
        r = getrandbits(k)  # 0 <= r < 2**k
        while r >= n:
            r = getrandbits(k)
        return r


class Patched(Random):
    # Above without the local
    def _randbelow(self, n):
        if not n:
            return 0
        k = n.bit_length()  # don't use (n-1) here because n can be 1
        r = self.getrandbits(k)  # 0 <= r < 2**k
        while r >= n:
            r = self.getrandbits(k)
        return r


class MainWithoutGetrandbits(Random):
    # Current main _randbelow_without_getrandbits()
    def _randbelow(self, n, maxsize=1<<BPF):
        random = self.random
        if n >= maxsize:
            _warn("Underlying random() generator does not supply \n"
                "enough bits to choose from a population range this large.\n"
                "To remove the range limitation, add a getrandbits() method.")
            return _floor(random() * n)
        if n == 0:
            return 0
        rem = maxsize % n
        limit = (maxsize - rem) / maxsize   # int(limit * maxsize) % n == 0
        r = random()
        while r >= limit:
            r = random()
        return _floor(r * maxsize) % n


class PatchedWithoutGetrandbits(Random):
    # Above without the local
    def _randbelow(self, n, maxsize=1<<BPF):
        if n >= maxsize:
            _warn("Underlying random() generator does not supply \n"
                "enough bits to choose from a population range this large.\n"
                "To remove the range limitation, add a getrandbits() method.")
            return _floor(self.random() * n)
        if n == 0:
            return 0
        rem = maxsize % n
        limit = (maxsize - rem) / maxsize   # int(limit * maxsize) % n == 0
        r = self.random()
        while r >= limit:
            r = self.random()
        return _floor(r * maxsize) % n

PAIRS = [(Main, Patched)]
if INCLUDE_WITHOUT_GETRANDBITS:
    PAIRS.append((MainWithoutGetrandbits, PatchedWithoutGetrandbits))

RANDOM_OBJECTS = {cls: cls(12345) for pair in PAIRS for cls in pair}


_dummy_lists = {n: [None] * n for n in (1024, 2000, 2047)}

def try_sample(r, n, k):
    f = r.sample
    l = _dummy_lists[n]
    start = time()
    for i in range(50):
        f(l, k)
    return time() - start

def try_shuffle(r, n):
    f = r.shuffle
    l = _dummy_lists[n]
    start = time()
    f(l)
    return time() - start

def try_randrange(r, n):
    f = r.randrange
    start = time()
    for i in range(1000):
        f(n)
    return time() - start

TESTS = {
    'sample([None] * 1024, 50)': lambda r: try_sample(r, 1024, 50),
    'sample([None] * 2047, 50)': lambda r: try_sample(r, 2047, 50),
    'shuffle([None] * 2000)': lambda r: try_shuffle(r, 2000),
    'randrange(1)': lambda r: try_randrange(r, 1),
    'randrange(2)': lambda r: try_randrange(r, 2),
    'randrange(3)': lambda r: try_randrange(r, 3),
    'randrange(4)': lambda r: try_randrange(r, 4),
    'randrange(15)': lambda r: try_randrange(r, 15),
    'randrange(16)': lambda r: try_randrange(r, 16),
    'randrange(1023)': lambda r: try_randrange(r, 1023),
    'randrange(1024)': lambda r: try_randrange(r, 1024),
    'randrange(2000000000)': lambda r: try_randrange(r, 2000000000),
}


total_times = {(cls, test): 0 for cls in RANDOM_OBJECTS for test in TESTS}
trials = list(total_times) * RANDOM_TRIAL_COUNT
shuffle(trials)

for cls, test in trials:
    total_times[cls, test] += TESTS[test](RANDOM_OBJECTS[cls])


print('                                        Speedups')
print('                     Test  with_getrandbits  without_getrandbits')

for test in TESTS:
    print(test.rjust(25), end='')
    for main, patched in PAIRS:
        speedup = 1 - total_times[patched, test] / total_times[main, test]
        print(f'{100 * speedup:5.1f}%'.rjust(18), end='')
    print()