# CelerisLab/cuda/obs_layout.py
"""Packed ``obs`` buffer layout for force / torque / sensor telemetry.

Usage::
    Host code (e.g. ``ObjectManager``) and ``compiler_v2.generate_config`` must use
    the same ``dim`` and ``n_objects`` so allocation sizes match device macros.

Design::
    Segments are contiguous in this order:
      1) force   : ``max(n_objects, 1) * dim`` floats
      2) torque  : ``max(n_objects, 1) * (1 if dim==2 else dim)`` floats
      3) sensor  : ``max(n_objects, 1) * dim`` floats
"""

from typing import NamedTuple


class ObsLayout(NamedTuple):
    """Integer offsets and sizes in **float** elements (not bytes)."""

    n_slots: int
    torque_components: int
    slot_stride_floats: int
    force_floats: int
    torque_floats: int
    force0_floats: int
    torque0_floats: int
    sensor0_floats: int
    total_floats: int


def obs_layout(dim: int, n_objects: int) -> ObsLayout:
    """Compute packed ``obs`` indices for force, torque, and sensor segments.

    Args:
        dim: Spatial dimension (2 or 3); must match ``LBMConfig.dim``.
        n_objects: Number of registered objects; may be zero.

    Returns:
        :class:`ObsLayout` with segment sizes and float offsets.
    """
    n_slots = max(int(n_objects), 1)
    torque_components = 1 if int(dim) == 2 else int(dim)
    slot_stride_floats = n_slots * dim
    force_floats = slot_stride_floats
    torque_floats = n_slots * torque_components
    force0_floats = 0
    torque0_floats = force0_floats + force_floats
    sensor0_floats = torque0_floats + torque_floats
    total_floats = sensor0_floats + slot_stride_floats
    return ObsLayout(
        n_slots=n_slots,
        torque_components=torque_components,
        slot_stride_floats=slot_stride_floats,
        force_floats=force_floats,
        torque_floats=torque_floats,
        force0_floats=force0_floats,
        torque0_floats=torque0_floats,
        sensor0_floats=sensor0_floats,
        total_floats=total_floats,
    )
