# CelerisLab/body/coupling/soa_packer.py
"""
SoA packer -- transforms CutLink records into kernel-consumable SoA arrays.

Responsibilities:
    - Pack CutLink[] -> 8-column SoA for CurvedLinkSoA.assign_host()
    - Pack SensorCell[] -> 2-column SoA for SensorSoA.assign_host()
    - Document the host-kernel SoA contract in one concentrated location.

Design:
    This is the single host-side function that knows the kernel SoA layout.
    Geometry builders (circle, polygon, mesh) produce CutLink[] and call this
    function -- they do NOT know the SoA format.  If the kernel SoA layout
    changes in the future, only this file needs updating.
"""

from __future__ import annotations

from typing import List, Tuple

import numpy as np

from ..geometry.base import CutLink, SensorCell


# Convenience aliases kept for backward compat during refactoring.
# Geometry callers now produce CutLink[] which is packed here.
# The constants remain for the fallback classification used during
# cut-link enumeration in circle.py.
FALLBACK_CLASS_BOUZIDI = np.uint8(0)
FALLBACK_CLASS_HALFWAY_BOUNCEBACK = np.uint8(1)


def pack_cut_links_to_soa(
    links: List[CutLink],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray,
           np.ndarray, np.ndarray, np.ndarray,
           np.ndarray, np.ndarray]:
    """Pack CutLink list into 8-column SoA arrays for CurvedLinkSoA.

    Output order matches CurvedLinkSoA.assign_host() signature::

        fluid_idx (uint32), dir (uint8), q (float32),
        body_id (int32), rx (float32), ry (float32),
        rz (float32), fallback (uint8)

    This is the single host-side SoA packing point.
    All geometry builders (circle, polygon, mesh) call this same
    function, isolating the kernel SoA layout from geometry logic.

    Semantic contract for each column (reference for host and kernel readers):

        fluid_idx
            Linear index of the fluid cell whose lattice link crosses into solid.
            This cell receives the Bouzidi-interpolated population on the next
            pull step.
        dir
            Lattice direction from fluid_idx toward the solid neighbor (1..NQ-1).
        q
            Fractional distance fluid -> hit_point / fluid -> neighbor, in (0, 1].
            Used by the kernel for Bouzidi interpolation weighting.
        body_id
            Object id used to look up wall velocity from the action buffer
            (action[body_id * slot + slot - 1] = omega).
        rx, ry, rz
            Lever arm from body center to wall hit point (hit_point - center),
            used by the kernel for torque accumulation:
            torque_z = rx * fy - ry * fx (2D).
        fallback
            0 = Bouzidi linear interpolation is legal (donor cell is fluid and
            in-domain).  1 = donor is invalid; use half-way bounce-back
            fallback [Gin08b].

    Args:
        links: List of CutLink records from any geometry builder.

    Returns:
        8-element tuple of numpy arrays matching CurvedLinkSoA.assign_host().
    """
    n = len(links)
    if n == 0:
        return (
            np.zeros(0, dtype=np.uint32),
            np.zeros(0, dtype=np.uint8),
            np.zeros(0, dtype=np.float32),
            np.zeros(0, dtype=np.int32),
            np.zeros(0, dtype=np.float32),
            np.zeros(0, dtype=np.float32),
            np.zeros(0, dtype=np.float32),
            np.zeros(0, dtype=np.uint8),
        )

    fluid_idx = np.empty(n, dtype=np.uint32)
    directions = np.empty(n, dtype=np.uint8)
    q_vals = np.empty(n, dtype=np.float32)
    body_ids = np.empty(n, dtype=np.int32)
    rx_vals = np.empty(n, dtype=np.float32)
    ry_vals = np.empty(n, dtype=np.float32)
    rz_vals = np.empty(n, dtype=np.float32)
    fallback_vals = np.empty(n, dtype=np.uint8)

    for i, cl in enumerate(links):
        fluid_idx[i] = np.uint32(cl.fluid_idx)
        directions[i] = np.uint8(cl.dir)
        q_vals[i] = np.float32(cl.q)
        body_ids[i] = np.int32(cl.body_id)
        rx_vals[i] = np.float32(cl.rx)
        ry_vals[i] = np.float32(cl.ry)
        rz_vals[i] = np.float32(cl.rz)
        fallback_vals[i] = np.uint8(cl.fallback)

    return (fluid_idx, directions, q_vals, body_ids,
            rx_vals, ry_vals, rz_vals, fallback_vals)


def pack_sensor_to_soa(
    cells: List[SensorCell],
) -> Tuple[np.ndarray, np.ndarray]:
    """Pack SensorCell list into 2-column SoA arrays for SensorSoA.

    Returns:
        cells (uint32), obj_ids (int32) -- empty arrays if *cells* is empty.
    """
    n = len(cells)
    if n == 0:
        return np.zeros(0, dtype=np.uint32), np.zeros(0, dtype=np.int32)
    idx_arr = np.empty(n, dtype=np.uint32)
    id_arr = np.empty(n, dtype=np.int32)
    for i, sc in enumerate(cells):
        idx_arr[i] = np.uint32(sc.idx)
        id_arr[i] = np.int32(sc.body_id)
    return idx_arr, id_arr


def pack_force_region_to_soa(
    cells: List[SensorCell],
) -> Tuple[np.ndarray, np.ndarray]:
    """Pack SensorCell list into 2-column SoA for ForceRegionSoA.

    Reuses SensorCell records since the footprint enumeration is identical
    to sensor cells.  Returns:
        cells (uint32), obj_ids (int32) -- empty arrays if *cells* is empty.
    """
    n = len(cells)
    if n == 0:
        return np.zeros(0, dtype=np.uint32), np.zeros(0, dtype=np.int32)
    idx_arr = np.empty(n, dtype=np.uint32)
    id_arr = np.empty(n, dtype=np.int32)
    for i, sc in enumerate(cells):
        idx_arr[i] = np.uint32(sc.idx)
        id_arr[i] = np.int32(sc.body_id)
    return idx_arr, id_arr
