# ADOBE CONFIDENTIAL
# Copyright 2025 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.

from __future__ import annotations
from typing import List, Tuple, Optional
import random

def stratified_random(n: int=64, seed: Optional[int]=None, shuffle_blocks: bool=True) -> List[int]:
    if n <= 0 or n & n - 1 != 0:
        raise ValueError('n must be a positive power of two (e.g., 64)')
    rng = random.Random(seed)
    occupied = [[False] * n for _ in range(n)]
    seq: List[int] = []
    blocks: List[Tuple[int, int, int]] = [(0, 0, n)]

    def block_has_sample(x0: int, y0: int, size: int) -> bool:
        for yy in range(y0, y0 + size):
            row = occupied[yy]
            for xx in range(x0, x0 + size):
                if row[xx]:
                    return True
        return False

    def place_random_in_block(x0: int, y0: int, size: int):
        x = rng.randrange(x0, x0 + size)
        y = rng.randrange(y0, y0 + size)
        attempts = 0
        while occupied[y][x]:
            x = rng.randrange(x0, x0 + size)
            y = rng.randrange(y0, y0 + size)
            attempts += 1
            if attempts > 10000:
                raise RuntimeError('Too many attempts to place a sample; logic error?')
        occupied[y][x] = True
        seq.append(y * n + x)
    size = n
    while size > 1:
        half = size // 2
        children: List[Tuple[int, int, int]] = []
        for (x0, y0, s) in blocks:
            assert s == size
            children.extend([(x0, y0, half), (x0 + half, y0, half), (x0, y0 + half, half), (x0 + half, y0 + half, half)])
        if shuffle_blocks:
            rng.shuffle(children)
        for (x0, y0, s) in children:
            if not block_has_sample(x0, y0, s):
                place_random_in_block(x0, y0, s)
        blocks = children
        size = half
    remaining: List[int] = []
    for y in range(n):
        for x in range(n):
            if not occupied[y][x]:
                remaining.append(y * n + x)
    rng.shuffle(remaining)
    seq.extend(remaining)
    assert len(seq) == n * n, (len(seq), n * n)
    return seq