"""Index construction for blocks and nested sub-blocks."""
from __future__ import annotations
import numpy as np


def equal_blocks(N: int, K: int):
    """Contiguous equal blocks; requires K | N. Returns array (K, m) of indices."""
    if N % K:
        raise ValueError(f"K={K} does not divide N={N}")
    m = N // K
    return np.arange(N).reshape(K, m)


def random_blocks(N: int, K: int, rng: np.random.Generator):
    """Random equal blocks (uniform permutation), (K, m)."""
    return rng.permutation(N).reshape(K, N // K)


def unequal_blocks(sizes):
    """Contiguous blocks with the given sizes -> list of index arrays."""
    sizes = [int(s) for s in sizes]
    out, start = [], 0
    for s in sizes:
        out.append(np.arange(start, start + s)); start += s
    return out


def nested_sub_blocks(block: np.ndarray, s: int):
    """Split one block (length m) into s equal contiguous sub-blocks: (s, m/s)."""
    m = block.shape[-1]
    if m % s:
        raise ValueError(f"split multiplier {s} does not divide block size {m}")
    return block.reshape(block.shape[:-1] + (s, m // s))


def check_splits(m: int, splits):
    for s in splits:
        if m % s:
            raise ValueError(f"block size m={m} not divisible by split multiplier s={s}")
    return True


def skewed_sizes(N: int, K: int, ratio: float):
    """Unequal block sizes: sizes grow geometrically so that max/min ~= ratio; sum = N (integer-rounded)."""
    g = np.geomspace(1.0, ratio, K)
    sizes = np.floor(g / g.sum() * N).astype(int)
    sizes[-1] += N - sizes.sum()
    return sizes