import math
import random
from multiprocessing import Pool
from typing import Optional, Iterable, Tuple

def gen_nums(n: int) -> list[int]:
    assert n > 0
    nums = list(range(2 ** n))
    random.shuffle(nums)
    return nums

def chunks(n: int, p: int) -> Iterable[Tuple[int, int]]:
    base, rem = divmod(n, p)
    start = 0
    for i in range(p):
        size = base + (1 if i < rem else 0)
        end = start + size
        yield start, end
        start = end

def t_find(args: Tuple[int, int, list[int], int]) -> Optional[int]:
    s, e, haystack, needle = args
    for i in range(s, e):
        if haystack[i] == needle:
            return i
    return None

def find(needle: int, haystack: list[int]) -> Optional[int]:
    size = len(haystack)
    assert size >= 1

    procs = int(math.log2(size))
    ranges = list(chunks(size, procs))
    with Pool(processes=procs) as pool:
        res = pool.map(t_find, [(s, e, haystack, needle) for (s, e) in ranges])

    for idx in res:
        if idx is not None:
            return idx
    return None

def fuzz_test(trials: int = 200, max_n: int = 16, seed: Optional[int] = None) -> None:
    if seed is not None:
        random.seed(seed)

    for t in range(1, trials + 1):
        n = random.randint(1, max_n)
        arr = gen_nums(n)

        idx_true = random.randrange(len(arr))
        needle = arr[idx_true]
        expected = idx_true

        got = find(needle, arr)
        assert got == expected, f"[t={t}] Expected {expected}, got {got}; n={n}"

    print(f"Fuzz OK: {trials} trials, max_n={max_n}")

fuzz_test(trials=1000, max_n=16)