Source code for simvx.graphics.renderer.fog_volume_pack

"""FogVolume3D std430 SSBO packing: backend-agnostic, numpy only.

The per-volume wire/SSBO layout is shared by the desktop Vulkan
``volumetric_fog_pass`` (which uploads it to a GPU buffer) and the web export's
``scene3d_serializer`` (which streams it to the browser WebGPU runtime). Keeping
the packing here, free of any ``vulkan`` import, lets the web export bundle it
into Pyodide without dragging in the desktop graphics stack. Both backends march
the identical bytes so the same scene fogs identically.
"""

from typing import Any

import numpy as np

__all__ = ["FOG_VOLUME_STRIDE", "MAX_FOG_VOLUMES", "STEP_COUNT", "build_volume_ssbo"]

# FogVolume SSBO stride: mat4 inv_transform(64) + vec4 albedo(16)
#   + vec4 params(16) + vec4 extra(16) = 112 bytes (std430, vec4-aligned).
FOG_VOLUME_STRIDE = 112
# Cap the per-frame volume SSBO. The buffer is allocated once at this size; the
# march loop only iterates the live count uploaded each frame.
MAX_FOG_VOLUMES = 64

# Ray-march step count (matches the web shader's fixed 32).
STEP_COUNT = 32




def build_volume_ssbo(volumes: list[Any]) -> tuple[np.ndarray, int]:
    """Pack ``FogVolume3D`` nodes into a flat std430 byte array.

    Returns ``(bytes_array, count)`` where ``count`` is clamped to
    ``MAX_FOG_VOLUMES``. Each record is::

        mat4 inv_transform   (world → unit-local, column-major for GLSL/WGSL)
        vec4 albedo          (rgb, a unused)
        vec4 params          (density, shape, edge_falloff, height_falloff)
        vec4 extra           (priority, 0, 0, 0)

    ``inv_transform`` folds the node's world transform with a per-axis scale of
    ``size * 0.5`` so the shape test in the shader is against a unit box / unit
    sphere / unit cylinder regardless of the volume's placement.
    """
    from simvx.core.math.matrices import mat4_from_trs

    n = min(len(volumes), MAX_FOG_VOLUMES)
    buf = np.zeros(max(1, n) * FOG_VOLUME_STRIDE, dtype=np.uint8)
    f32 = buf.view(np.float32)
    for i in range(n):
        v = volumes[i]
        size = v.size
        ws = v.world_scale
        half = np.array(
            [max(float(size[0]) * 0.5 * float(ws[0]), 1e-4),
             max(float(size[1]) * 0.5 * float(ws[1]), 1e-4),
             max(float(size[2]) * 0.5 * float(ws[2]), 1e-4)],
            dtype=np.float32,
        )
        # World transform of the node (row-major), with the per-axis scale set
        # to half-extents so the shape test in the shader is against a unit
        # box / sphere / cylinder regardless of placement.
        model = mat4_from_trs(
            v.world_position, v.world_rotation, half,
        ).astype(np.float32)
        try:
            inv = np.linalg.inv(model).astype(np.float32)
        except np.linalg.LinAlgError:
            inv = np.eye(4, dtype=np.float32)

        base = i * FOG_VOLUME_STRIDE // 4
        # GLSL/WGSL std430 mat4 is column-major; numpy is row-major → transpose.
        f32[base:base + 16] = inv.T.ravel()
        alb = v.albedo
        f32[base + 16:base + 20] = [float(alb[0]), float(alb[1]), float(alb[2]),
                                    float(alb[3]) if len(alb) > 3 else 1.0]
        f32[base + 20:base + 24] = [
            v.effective_density,
            float(int(v.shape)),
            float(v.falloff),
            float(v.height_falloff),
        ]
        f32[base + 24] = float(v.priority)
    return buf, n