from dataclasses import dataclass


@dataclass(frozen=True)
class OffloadSelection:
    enabled: bool
    layer_indices: tuple[int, ...]


def select_offload_layers(
    *,
    num_layers: int,
    group_size: int,
    num_in_group: int,
) -> OffloadSelection:
    if group_size < 0:
        raise ValueError("offload_group_size must be >= 0")
    if num_in_group < 0:
        raise ValueError("offload_num_in_group must be >= 0")
    if group_size == 0 or num_in_group == 0 or num_layers <= 0:
        return OffloadSelection(enabled=False, layer_indices=())
    if num_in_group > group_size:
        raise ValueError(
            f"offload_num_in_group ({num_in_group}) must be <= "
            f"offload_group_size ({group_size})"
        )

    cutoff = group_size - num_in_group
    indices = tuple(
        layer_idx
        for layer_idx in range(num_layers)
        if layer_idx % group_size >= cutoff
    )
    return OffloadSelection(enabled=bool(indices), layer_indices=indices)
