# CelerisLab/body/manager.py
"""
ObjectManager -- batch management of SimObjects.

Responsibilities:
  - Add / remove / query objects (delegated to ``BodyRegistry``)
  - Pending edit state for runtime body topology changes
  - Build merged flag masks and compact cut-link / sensor lists
  - Allocate packed telemetry ``obs_gpu`` + pagelocked mirror ``obs_pinned``
  - Sync geometry to :class:`~CelerisLab.lbm.field.LBMField`

Design::
  Packed ``obs`` layout matches ``generate_config(..., n_objects=count)`` and
  ``config_obs.h`` macros. Host sizing uses :func:`~CelerisLab.cuda.compiler_v2.obs_layout`.

  Pending edits are lightweight: only the "intent" (add / remove) is stored.
  All derived quantities (flags, compact lists, GPU buffers) are built at
  ``sync_bodies()`` time from the formal registry + pending edits.
"""

from __future__ import annotations

from typing import Dict, List, NamedTuple, Optional, TYPE_CHECKING

import numpy as np
import pycuda.driver as cuda

from .objects import SimObject
from .registry import BodyRegistry
from .preprocess.flag_overlay import merge_flag_masks
from ..cuda.compiler_v2 import obs_layout
from ..cuda.obs_layout import ObsLayout
from ..lbm.descriptors import (
    D2Q9_EX,
    D2Q9_EY,
    D3Q19_EX,
    D3Q19_EY,
    D3Q19_EZ,
)

if TYPE_CHECKING:
    from .sync_plan import BodySyncPlan


class BodyTelemetry(NamedTuple):
    """Unified telemetry for one body, read from the pinned obs buffer.

    After ``run(sync_obs=True)``, all fields are populated from GPU results.
    ``sensor`` is a dim-length array — zeros for non-sensor bodies.
    """
    force: np.ndarray
    torque: np.ndarray
    sensor: np.ndarray


class ObjectManager:
    """Central registry for all simulation objects.

    Frozen telemetry layout::
        ``obs`` packs force, torque, then sensor segments.
        Segment offsets and sizes come from :func:`obs_layout(dim, count)`.
        Device macros ``OBS_*`` in ``config_obs.h`` must match ``obs_layout(dim, count)``.
    """

    def __init__(self, nx: int, ny: int, nz: int, nq: int, cfg=None):
        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.nq = nq
        self.cfg = cfg
        self._registry = BodyRegistry()

        self.action_gpu: Optional[cuda.DeviceAllocation] = None
        self.obs_gpu: Optional[cuda.DeviceAllocation] = None
        self._action_nbytes: int = 0
        self._obs_alloc_nbytes: int = 0

        self.action = np.zeros(0, dtype=np.float32)
        self.obs_pinned: Optional[np.ndarray] = None
        self.obs_force_view: Optional[np.ndarray] = None
        self.obs_torque_view: Optional[np.ndarray] = None
        self.obs_sensor_view: Optional[np.ndarray] = None
        self.sensor_cell_counts = np.zeros(0, dtype=np.int32)
        self.slot_stride_floats: int = 0
        self.torque_stride_floats: int = 0
        self.torque_components: int = 0
        self.torque0_floats: int = 0
        self.sensor0_floats: int = 0
        self.obs_total_floats: int = 0
        self.obs_nbytes: int = 0

        self._telemetry_field: Optional[object] = None

        # -- Pending edit state (runtime body topology sync) -------------------
        self._edit_active: bool = False
        self._pending_add: List[SimObject] = []
        self._pending_remove: set[int] = set()

        # Ensure host action buffer is non-empty so CUDA alloc never sees 0 bytes
        # when bodies.count == 0 (matches _resize_buffers sizing for n=0).
        self._resize_buffers()

    # -- Registry delegation -------------------------------------------------
    def add(self, obj: SimObject) -> int:
        """Register an object and return its id."""
        body_id = self._registry.add(obj)
        self._resize_buffers()
        return body_id

    def remove(self, obj_id: int):
        self._registry.remove(obj_id)
        self._resize_buffers()

    def get(self, obj_id: int) -> SimObject:
        return self._registry.get(obj_id)

    @property
    def objects(self) -> List[SimObject]:
        return self._registry.objects

    @property
    def count(self) -> int:
        return self._registry.count

    # -- Buffer sizing -------------------------------------------------------
    def _resize_buffers(self):
        n = self.count
        dim = self.cfg.dim if self.cfg else 2
        action_slot_floats = 3 * dim
        new_action = np.zeros(max(n * action_slot_floats, 1), dtype=np.float32)
        if self.action.size > 0:
            copy_n = min(self.action.size, new_action.size)
            new_action[:copy_n] = self.action[:copy_n]
        self.action = new_action

        new_sensor_counts = np.zeros(max(n, 1), dtype=np.int32)
        if self.sensor_cell_counts.size > 0:
            copy_n = min(self.sensor_cell_counts.size, new_sensor_counts.size)
            new_sensor_counts[:copy_n] = self.sensor_cell_counts[:copy_n]
        self.sensor_cell_counts = new_sensor_counts

    # -- Pending edit state (runtime body topology sync) ---------------------
    def stage_add(self, obj: SimObject) -> None:
        """Stage an object for addition at the next ``sync_bodies()``.

        The object is not added to the formal registry until ``commit_pending()``
        is called.  This implicitly activates the edit window.
        """
        self._edit_active = True
        self._pending_add.append(obj)

    def stage_remove(self, obj_id: int) -> None:
        """Stage removal of a formal object at the next ``sync_bodies()``.

        Args:
            obj_id: Id of a currently registered formal object.

        Raises:
            IndexError: If *obj_id* is not in the formal registry.
        """
        self._validate_body_id(obj_id)
        self._edit_active = True
        self._pending_remove.add(obj_id)

    def has_pending_edit(self) -> bool:
        """Return True if there are uncommitted staged edits."""
        return self._edit_active and bool(self._pending_add or self._pending_remove)

    def clear_pending_edits(self) -> None:
        """Discard all staged edits without applying them."""
        self._pending_add.clear()
        self._pending_remove.clear()
        self._edit_active = False

    # -- Sync plan construction (runtime body topology sync) -----------------
    def build_next_objects(self) -> List[SimObject]:
        """Merge formal objects with pending edits into a new object list.

        Formal objects whose id is in ``_pending_remove`` are excluded.
        Objects in ``_pending_add`` are appended.  All returned objects
        receive consecutive ``obj_id = 0..n-1``.

        Returns:
            New list of SimObject copies with reassigned ids.
        """
        def _copy(obj, oid):
            new = SimObject(obj_id=oid, geometry=obj.geometry,
                            center=obj.center, radius=obj.radius,
                            is_sensor=obj.is_sensor,
                            is_force_region=obj.is_force_region)
            new.state = obj.state
            new.control = obj.control
            return new

        result: List[SimObject] = []
        next_id = 0
        for obj in self._registry.objects:
            if obj.obj_id in self._pending_remove:
                continue
            result.append(_copy(obj, next_id))
            next_id += 1
        for obj in self._pending_add:
            result.append(_copy(obj, next_id))
            next_id += 1
        return result

    def build_sync_plan(self, field) -> 'BodySyncPlan':
        """Construct a full topology transition plan from pending edits.

        Args:
            field: ``LBMField`` instance (used for old flags and channel base).

        Returns:
            :class:`BodySyncPlan` describing the transition.
        """
        from .sync_plan import BodySyncPlan
        from ..lbm.descriptors import FLUID, OBSTACLE

        old_flags = field.flag.copy()

        next_objects = self.build_next_objects()
        next_count = len(next_objects)

        # Build new flags from clean channel base + new object set
        base_flags = field.build_channel_flags()
        next_flags = self.build_flags_for(
            next_objects, base_flags,
            nx=self.nx, ny=self.ny, nz=self.nz,
        )

        # Build compact lists for new object set
        new_sensor_counts = np.zeros(max(next_count, 1), dtype=np.int32)
        (
            cl_fluid_idx, cl_dir, cl_q, cl_body_id,
            cl_rx, cl_ry, cl_rz, cl_fallback_class,
            sensor_cells, sensor_obj_id,
            fr_cells, fr_obj_id,
        ) = self.build_compact_lists_for(
            next_objects, domain_flags=next_flags,
            sensor_cell_counts_out=new_sensor_counts,
        )

        # Compute change masks
        old_is_solid = (old_flags & 0x0002) != 0
        new_is_solid = (next_flags & 0x0002) != 0
        old_is_obstacle = (old_flags & OBSTACLE) != 0
        new_is_fluid = (next_flags & FLUID) != 0

        added_solid_mask = (~old_is_solid) & new_is_solid
        released_fluid_mask = old_is_obstacle & new_is_fluid

        return BodySyncPlan(
            next_objects=next_objects,
            next_count=next_count,
            next_flags=next_flags,
            curved_host=(cl_fluid_idx, cl_dir, cl_q, cl_body_id,
                         cl_rx, cl_ry, cl_rz, cl_fallback_class),
            sensor_host=(sensor_cells, sensor_obj_id),
            force_region_host=(fr_cells, fr_obj_id),
            added_solid_mask=added_solid_mask,
            released_fluid_mask=released_fluid_mask,
            new_sensor_counts=new_sensor_counts,
        )

    def commit_pending(self, next_objects: List[SimObject],
                       sensor_counts: np.ndarray) -> None:
        """Replace the formal registry with *next_objects* and clear pending edits.

        Args:
            next_objects: New formal object list with ``obj_id = 0..n-1``.
            sensor_counts: New per-object sensor cell counts.
        """
        self._registry.clear()
        for obj in next_objects:
            self._registry.add(obj)
        self.sensor_cell_counts = sensor_counts.copy()
        self._resize_buffers()
        self.clear_pending_edits()

    def apply_sync_plan(self, field, plan) -> None:
        """Apply topology data from a ``BodySyncPlan`` to the field.

        Sets ``field.flag``, assigns compact list host arrays, uploads
        compact lists, and refreshes action buffer from object state.

        Args:
            field: LBMField instance.
            plan: BodySyncPlan from ``build_sync_plan()``.
        """
        field.flag = plan.next_flags
        field.upload_flags()

        cl_host = plan.curved_host
        field.curved.assign_host(*cl_host)
        s_cells, s_ids = plan.sensor_host
        field.sensors.assign_host(s_cells, s_ids)
        fr_cells, fr_ids = plan.force_region_host
        field.force_regions.assign_host(fr_cells, fr_ids)

        field.upload_compact_lists()
        field.update_params(n_objects=plan.next_count)

    # -- Flag composition ----------------------------------------------------
    def build_flags(self, base_flags: np.ndarray) -> np.ndarray:
        """Merge formal object flag masks onto a clean domain base."""
        return self.build_flags_for(
            self._registry.objects, base_flags,
            nx=self.nx, ny=self.ny, nz=self.nz,
        )

    @staticmethod
    def build_flags_for(objects: List[SimObject],
                        base_flags: np.ndarray,
                        *, nx: int, ny: int, nz: int) -> np.ndarray:
        """Merge object flag masks onto a clean domain base.

        Args:
            objects: List of SimObject instances.
            base_flags: Fresh channel-layout flag array.
            nx, ny, nz: Grid dimensions.

        Returns:
            New uint16 array with object overlays applied.
        """
        masks = [obj.get_flag_mask(nx, ny, nz) for obj in objects]
        return merge_flag_masks(base_flags, masks)

    # -- Compact list building -----------------------------------------------
    def build_compact_lists(self, domain_flags: np.ndarray | None = None):
        """Build cut-link SoA columns, sensor lists, and force-region lists
        from the formal registry objects.
        """
        return self.build_compact_lists_for(
            self._registry.objects, domain_flags,
            sensor_cell_counts_out=self.sensor_cell_counts,
        )

    def build_compact_lists_for(
        self,
        objects: List[SimObject],
        domain_flags: np.ndarray | None = None,
        *,
        sensor_cell_counts_out: np.ndarray | None = None,
    ):
        """Build cut-link SoA columns, sensor lists, and force-region lists.

        Args:
            objects: List of SimObject instances to process.
            domain_flags: Optional uint16 flags for donor-cell classification.
            sensor_cell_counts_out: Int array to fill with per-object sensor
                cell counts.  If provided, ``fill(0)`` is called first.

        Returns:
            cl_fluid_idx, cl_dir, cl_q, cl_body_id, cl_rx, cl_ry, cl_rz,
            cl_fallback_class,
            sensor_cells, sensor_obj_id,
            fr_cells, fr_obj_id
        """
        cl_idx, cl_dir, cl_q, cl_bid = [], [], [], []
        cl_rx, cl_ry, cl_rz, cl_fallback = [], [], [], []
        s_cells, s_ids = [], []
        fr_cells, fr_ids = [], []

        if sensor_cell_counts_out is not None:
            sensor_cell_counts_out.fill(0)

        ez = None
        if self.cfg and self.cfg.is_d3q19:
            ex, ey, ez = D3Q19_EX, D3Q19_EY, D3Q19_EZ
        else:
            ex, ey = D2Q9_EX, D2Q9_EY

        for obj in objects:
            if (not obj.is_sensor) and (not obj.is_force_region) and hasattr(obj, 'get_curved_list'):
                (
                    fluid_idx, dirs, q_vals, body_ids,
                    rx_vals, ry_vals, rz_vals, fallback_vals
                ) = obj.get_curved_list(
                    self.nx, self.ny, self.nq,
                    nz=self.nz, ex=ex, ey=ey, ez=ez,
                    domain_flags=domain_flags,
                )
                if len(fluid_idx) > 0:
                    cl_idx.append(fluid_idx)
                    cl_dir.append(dirs)
                    cl_q.append(q_vals)
                    cl_bid.append(body_ids)
                    cl_rx.append(rx_vals)
                    cl_ry.append(ry_vals)
                    cl_rz.append(rz_vals)
                    cl_fallback.append(fallback_vals)
            if obj.is_sensor and hasattr(obj, 'get_sensor_list'):
                cells, ids = obj.get_sensor_list(self.nx, self.ny, self.nz)
                if len(cells) > 0:
                    s_cells.append(cells)
                    s_ids.append(ids)
                    if (sensor_cell_counts_out is not None
                            and 0 <= obj.obj_id < sensor_cell_counts_out.size):
                        sensor_cell_counts_out[obj.obj_id] = int(len(cells))
            if obj.is_force_region and hasattr(obj, 'get_sensor_list'):
                cells, ids = obj.get_sensor_list(self.nx, self.ny, self.nz)
                if len(cells) > 0:
                    fr_cells.append(cells)
                    fr_ids.append(ids)

        r_idx = np.concatenate(cl_idx) if cl_idx else np.zeros(0, dtype=np.uint32)
        r_dir = np.concatenate(cl_dir) if cl_dir else np.zeros(0, dtype=np.uint8)
        r_q = np.concatenate(cl_q) if cl_q else np.zeros(0, dtype=np.float32)
        r_bid = np.concatenate(cl_bid) if cl_bid else np.zeros(0, dtype=np.int32)
        r_rx = np.concatenate(cl_rx) if cl_rx else np.zeros(0, dtype=np.float32)
        r_ry = np.concatenate(cl_ry) if cl_ry else np.zeros(0, dtype=np.float32)
        r_rz = np.concatenate(cl_rz) if cl_rz else np.zeros(0, dtype=np.float32)
        r_fallback = (
            np.concatenate(cl_fallback) if cl_fallback
            else np.zeros(0, dtype=np.uint8)
        )
        sensor_cells = np.concatenate(s_cells) if s_cells else np.zeros(0, dtype=np.uint32)
        sensor_obj_id = np.concatenate(s_ids) if s_ids else np.zeros(0, dtype=np.int32)
        r_fr_cells = np.concatenate(fr_cells) if fr_cells else np.zeros(0, dtype=np.uint32)
        r_fr_ids = np.concatenate(fr_ids) if fr_ids else np.zeros(0, dtype=np.int32)

        return (
            r_idx, r_dir, r_q, r_bid,
            r_rx, r_ry, r_rz, r_fallback,
            sensor_cells, sensor_obj_id,
            r_fr_cells, r_fr_ids,
        )

    # -- GPU sync ------------------------------------------------------------
    def sync_to_gpu(self, field, *, rebuild_flags: bool = True):
        """Upload compact lists, action buffer, and packed ``obs`` GPU buffer.

        Args:
            rebuild_flags: Rebuild ``field.flag`` from the clean channel base and
                upload it before constructing compact lists. Pass ``False`` when
                the caller has already prepared and uploaded the final flag field
                for this geometry, such as during initialization.
        """
        self._telemetry_field = field

        if rebuild_flags:
            field.flag = self.build_flags(field.build_channel_flags())
            field.upload_flags()

        (
            cl_fluid_idx, cl_dir, cl_q, cl_body_id,
            cl_rx, cl_ry, cl_rz, cl_fallback_class,
            sensor_cells, sensor_obj_id,
            fr_cells, fr_obj_id,
        ) = self.build_compact_lists(domain_flags=field.flag)

        field.curved.assign_host(
            cl_fluid_idx, cl_dir, cl_q, cl_body_id,
            cl_rx, cl_ry, cl_rz, cl_fallback_class,
        )
        field.sensors.assign_host(sensor_cells, sensor_obj_id)
        field.force_regions.assign_host(fr_cells, fr_obj_id)

        field.upload_compact_lists()
        field.update_params(n_objects=self.count)

        self._refresh_action_from_objects()
        self._allocate_packed_telemetry()
        assert self.obs_pinned is not None
        self.obs_pinned.fill(0)
        cuda.memcpy_htod(self.obs_gpu, self.obs_pinned)

    def _allocate_packed_telemetry(self) -> None:
        """Resize ``action_gpu`` / ``obs_gpu`` / ``obs_pinned`` when layout changes."""
        dim = self.cfg.dim if self.cfg else 2
        lay = obs_layout(dim, self.count)
        self._apply_obs_layout(lay)

        action_nbytes = int(self.action.nbytes)
        if self.action_gpu is None or self._action_nbytes != action_nbytes:
            if self.action_gpu is not None:
                self.action_gpu.free()
                self.action_gpu = None
            self.action_gpu = cuda.mem_alloc(action_nbytes)
            self._action_nbytes = action_nbytes
            # Initialise with host action (zeros until first run() upload).
            cuda.memcpy_htod(self.action_gpu, self.action)

        if self.obs_gpu is None or self._obs_alloc_nbytes != self.obs_nbytes:
            if self.obs_gpu is not None:
                self.obs_gpu.free()
                self.obs_gpu = None
            self.obs_gpu = cuda.mem_alloc(self.obs_nbytes)
            self._obs_alloc_nbytes = self.obs_nbytes
            self.obs_pinned = cuda.pagelocked_empty(lay.total_floats, np.float32)
            self.obs_pinned.fill(0)
            n_slots = max(self.count, 1)
            self.obs_force_view = self.obs_pinned[
                lay.force0_floats: lay.force0_floats + lay.force_floats
            ].reshape(
                n_slots, dim)
            self.obs_torque_view = self.obs_pinned[
                lay.torque0_floats: lay.torque0_floats + lay.torque_floats
            ].reshape(n_slots, lay.torque_components)
            self.obs_sensor_view = self.obs_pinned[
                lay.sensor0_floats: lay.sensor0_floats + lay.slot_stride_floats
            ].reshape(n_slots, dim)

    def _apply_obs_layout(self, lay: ObsLayout) -> None:
        self.slot_stride_floats = lay.slot_stride_floats
        self.torque_stride_floats = lay.torque_floats
        self.torque_components = lay.torque_components
        self.torque0_floats = lay.torque0_floats
        self.sensor0_floats = lay.sensor0_floats
        self.obs_total_floats = lay.total_floats
        self.obs_nbytes = lay.total_floats * 4

    def ensure_packed_buffers(self) -> None:
        """Allocate ``action_gpu`` and packed ``obs`` when there are zero objects.

        Call after ``generate_config(..., n_objects=0)`` so ``config_obs.h`` matches
        :func:`obs_layout(dim, 0)`.
        """
        if self.cfg is None:
            return
        self._allocate_packed_telemetry()
        assert self.obs_pinned is not None
        self.obs_pinned.fill(0)
        cuda.memcpy_htod(self.obs_gpu, self.obs_pinned)

    def zero_force_segment_async(self, stream: cuda.Stream):
        """Zero body telemetry (force + torque) of ``obs_gpu``."""
        n_floats = self.sensor0_floats
        cuda.memset_d32_async(self.obs_gpu, 0, n_floats, stream)

    def zero_sensor_segment_async(self, stream: cuda.Stream):
        """Zero the sensor segment (second stride-sized block of floats).

        This always issues a ``memset`` on the sensor sub-range. Call it only when
        the sensor kernel runs (e.g. ``field.n_sensor > 0``); the runner decides.
        """
        offset_bytes = self.sensor0_floats * 4
        ptr = int(self.obs_gpu) + offset_bytes
        cuda.memset_d32_async(ptr, 0, self.slot_stride_floats, stream)

    def download_obs_full_async(self, stream: cuda.Stream):
        """Enqueue full DTOH copy ``obs_gpu`` -> ``obs_pinned``."""
        assert self.obs_pinned is not None
        cuda.memcpy_dtoh_async(self.obs_pinned, self.obs_gpu, stream)

    def _upload_action_async(self, stream: cuda.Stream) -> None:
        """Async H2D of host action array to action_gpu."""
        if self.action_gpu is None or self.action.size == 0:
            return
        cuda.memcpy_htod_async(self.action_gpu, self.action, stream)

    @property
    def obs_n_slots(self) -> int:
        """At-least-one slot count ``max(count, 1)`` matching ``config_obs.h``."""
        return max(self.count, 1)

    @property
    def obs_slot_stride_floats(self) -> int:
        """Floats per telemetry segment (``obs_n_slots * DIM``)."""
        return self.slot_stride_floats

    @property
    def obs_sensor_base_floats(self) -> int:
        """Float index where the sensor segment begins."""
        return self.sensor0_floats

    @property
    def obs_torque_base_floats(self) -> int:
        """Float index where the torque segment begins."""
        return self.torque0_floats

    def read_force(self, body_id: int) -> np.ndarray:
        """Return the DIM-vector force on body ``body_id`` from ``obs_pinned``.

        Caller must have synchronised the CUDA stream before reading.
        """
        self._validate_body_id(body_id)
        assert self.cfg is not None
        assert self.obs_pinned is not None
        d = self.cfg.dim
        i0 = body_id * d
        return np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)

    def read_torque(self, body_id: int) -> np.ndarray:
        """Return torque vector for ``body_id`` from ``obs_pinned``."""
        self._validate_body_id(body_id)
        assert self.obs_pinned is not None
        i0 = self.torque0_floats + body_id * self.torque_components
        return np.array(
            self.obs_pinned[i0:i0 + self.torque_components], dtype=np.float32)

    def read_sensor(self, body_id: int, *, normalize: bool = True) -> np.ndarray:
        """Return sensor accumulation for ``body_id``.

        By default this returns the area-averaged value over the sensor footprint.
        Set ``normalize=False`` to get the raw sum accumulated by ``SensorKernel``.
        """
        self._validate_body_id(body_id)
        assert self.obs_pinned is not None
        assert self.cfg is not None
        d = self.cfg.dim
        i0 = self.sensor0_floats + body_id * d
        values = np.array(self.obs_pinned[i0:i0 + d], dtype=np.float32)
        if not normalize:
            return values
        count = int(self.sensor_cell_counts[body_id]) if body_id < self.sensor_cell_counts.size else 0
        if count <= 0:
            return values
        return values / np.float32(count)

    def read_body(self, body_id: int) -> BodyTelemetry:
        """Return unified telemetry for one body from the pinned obs buffer.

        The caller must ensure ``run(sync_obs=True)`` or an explicit
        ``download_obs_full_async + synchronize`` has completed.
        """
        force = self.read_force(body_id)
        torque = self.read_torque(body_id)
        sensor = self.read_sensor(body_id, normalize=True)
        return BodyTelemetry(force=force, torque=torque, sensor=sensor)

    def _obs_array(self) -> np.ndarray:
        """Return the full pinned obs buffer as a flat float32 array.

        Shape: ``(total_floats,)`` with layout ``[force|torque|sensor]``.
        Efficient for DRL batch reads.
        """
        assert self.obs_pinned is not None
        return self.obs_pinned.copy()

    def set_body_state(self, body_id: int, omega: float = 0.0) -> None:
        """Set runtime body state (host action array only, no H2D).

        The GPU action buffer is updated at the next ``run()`` call.
        """
        self._validate_body_id(body_id)
        dim = self.cfg.dim if self.cfg else 2
        slot = 3 * dim
        base = body_id * slot
        self.action[base + slot - 1] = np.float32(omega)

    def set_force_state(self, body_id: int, fx: float = 0.0, fy: float = 0.0) -> None:
        """Set force density on a force_region object (host action only, no H2D).

        The GPU action buffer is updated at the next ``run()`` call.
        """
        self._validate_body_id(body_id)
        dim = self.cfg.dim if self.cfg else 2
        slot = 3 * dim
        base = body_id * slot
        self.action[base + 0] = np.float32(fx)
        self.action[base + 1] = np.float32(fy)

    def _validate_body_id(self, body_id: int) -> None:
        if body_id < 0 or body_id >= self.count:
            raise IndexError(
                f"body_id {body_id} out of range for {self.count} objects")

    def _refresh_action_from_objects(self) -> None:
        """Populate action slots from object state using the runtime contract.

        Slot layout per object (DIM=2 in this example, slots derived from 3*DIM):
            [0] = vx   (= fx for force_region objects; 0 for bodies/sensors)
            [1] = vy   (= fy for force_region objects; 0 for bodies/sensors)
            [5] = omega (= body rotation for circle objects; 0 for sensors/force_region)

        Slot 0/1 are consumed by ForceRegionKernel for force_region objects.
        Slot -1 is consumed by CurvedBoundaryKernel for body rotation (action_omega).
        Sensors do not consume action slots.
        """
        if self.count == 0:
            return
        dim = self.cfg.dim if self.cfg else 2
        slot = 3 * dim
        self.action.fill(0)
        for obj in self._registry.objects:
            base = obj.obj_id * slot
            if base + slot > self.action.size:
                continue
            self.action[base + 0] = np.float32(getattr(obj.state, "vx", 0.0))
            self.action[base + 1] = np.float32(getattr(obj.state, "vy", 0.0))
            self.action[base + slot - 1] = np.float32(
                getattr(obj.state, "omega", 0.0))

    def fallback_link_count(self) -> int:
        """Return number of curved links marked for fallback bounce-back."""
        if self._telemetry_field is None:
            return 0
        curved = self._telemetry_field.curved
        if curved.count == 0:
            return 0
        return int(np.count_nonzero(curved.fallback_class[:curved.count] != 0))

    def low_q_link_count(self, threshold: float = 0.5) -> int:
        """Return count of links where q is below ``threshold``."""
        if self._telemetry_field is None:
            return 0
        curved = self._telemetry_field.curved
        if curved.count == 0:
            return 0
        return int(np.count_nonzero(
            curved.q[:curved.count] < np.float32(threshold)))

    # -- Placeholders for future use -----------------------------------------
    def update_states(self, dt: float):
        """Integrate object motion (placeholder)."""
        pass

    def detect_collisions(self):
        """Multi-body collision detection (placeholder)."""
        pass

    def exchange_forces(self, obs: np.ndarray):
        """Distribute fluid forces to objects (placeholder)."""
        pass
