"""Utilities for creating hierarchical orderings of sequences.

These functions create different patterns for processing sequences in a hierarchical manner,
enabling models to learn dependencies at multiple temporal scales.
"""


def create_interleaved_order_delay(N: int, K: int = 5, DELAY: int = 2, PAD: int = -1) -> list:
    """
    Creates an interleaved ordering of indices with delays between groups.

    This creates a hierarchical pattern where tokens are processed at multiple temporal scales,
    allowing the model to learn dependencies at different levels of granularity.

    Example with N=10, K=5, DELAY=2:
        Input: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        Groups: [0, 5], [1, 6], [2, 7], [3, 8], [4, 9]
        With delays: [0, 5], [_, _, 1, 6], [_, _, _, _, 2, 7], ...
        Interleaved: [0, _, _, _, _, 5, _, _, _, _, _, 1, _, _, _, _, 6, ...]

    Args:
        N: Total number of positions to interleave
        K: Number of groups to create (elements with same modulo K)
        DELAY: Number of PAD tokens to insert before each group (progressive delay)
        PAD: Value to use for padding tokens

    Returns:
        List of interleaved indices with padding, creating a hierarchical pattern
    """
    pos = list(range(N))
    mod_order = list(range(K))

    # Step 1: Split positions into K groups based on modulo
    # Group i contains all positions where pos % K == i
    # E.g., for K=5: Group 0 = [0, 5, 10, ...], Group 1 = [1, 6, 11, ...], etc.
    mods = []
    for i, mod in enumerate(mod_order):
        mod_pos = pos[mod::K]  # Extract every K-th element starting at mod

        # Add progressive delay: Group i gets i*DELAY padding tokens at the start
        # This creates a temporal hierarchy where later groups appear progressively later
        mod_pos = [PAD] * i * DELAY + mod_pos
        mods.append(mod_pos)

    # Step 2: Interleave the groups round-robin
    # At each position, take one element from each group (if available)
    # This creates patterns like: [g0[0], g1[0], g2[0], g3[0], g4[0], g0[1], g1[1], ...]
    max_len = max(len(mod) for mod in mods)
    final_order = []
    for i in range(max_len):
        for mod in mods:
            if len(mod) > i:
                final_order.append(mod[i])
            else:
                # Pad if this group has no more elements
                final_order.append(PAD)

    return final_order


def create_binary_tree_order(N: int) -> list:
    """
    Creates a hierarchical ordering using recursive binary partitioning over indices 0 to N-1.

    This performs breadth-first binary splits, visiting midpoints of progressively smaller
    ranges. Creates a coarse-to-fine hierarchy where larger-scale structure is processed
    before finer details.

    Example with N=7:
        [0, 1, 2, 3, 4, 5, 6]

        Depth 0: Split [0-7) → middle = 3
        Depth 1: Split [0-3) → middle = 1, Split [4-7) → middle = 5
        Depth 2: Split [0-1) → 0, Split [2-3) → 2, Split [4-5) → 4, Split [6-7) → 6

        Result: [3, 1, 5, 0, 2, 4, 6]

    Example with N=10:
        Result: [5, 2, 8, 1, 4, 7, 9, 0, 3, 6]
        Depth 0: [5]
        Depth 1: [2, 8] (middles of [0-5) and [6-10))
        Depth 2: [1, 4, 7, 9]
        Depth 3: [0, 3, 6]

    Args:
        N: Total number of positions to order

    Returns:
        List of indices ordered by hierarchical binary partitioning (breadth-first)
    """
    if N == 0:
        return []

    result = []
    queue = [(0, N)]  # Queue of (start, end) ranges

    while queue:
        start, end = queue.pop(0)
        if start >= end:
            continue

        mid = (start + end) // 2
        result.append(mid)

        # Add left and right subtrees to queue for breadth-first processing
        queue.append((start, mid))
        queue.append((mid + 1, end))

    return result