import torch
import itertools


def get_stratified_coords(ratio, shape):
    """
    Produce a list of approx. ``num_pix`` random coordinates, sampled from
    ``shape`` using startified sampling. Supports n-dimensional shapes.
    """
    # total_num = torch.prod(shape).to(torch.float32)
    # sample_num = total_num * ratio
    ratio = torch.as_tensor(ratio)
    ndim = len(shape)
    shape = torch.as_tensor(shape, dtype=torch.int32)
    box_size = int(torch.round(torch.sqrt(1. / ratio)))
    coords = []
    box_counts = torch.ceil(shape.float() / box_size).int()
    for steps in itertools.product(*[range(bc) for bc in box_counts]):
        steps = torch.as_tensor(steps, dtype=torch.int32)
        co = torch.randint(0, box_size, (ndim,)) + box_size * steps
        if torch.all(co < shape):
            coords.append(co)
    if not coords:
        raise ValueError(f'ratio {ratio:.1e} is too close to zero. Choose a higher value.')
    coords = torch.stack(coords)
    return coords