# CelerisLab/lbm/equilibrium.py
"""
Host-side equilibrium and macroscopic helpers for D2Q9. (D2Q9 only)

These functions are shared by LBMField (macroscopic computation) and
body DDF patching.  They operate on host NumPy arrays only and do not
require a CUDA module.

For the authoritative GPU-side equilibrium, see ``operators/equilibrium.cuh``.

Callers with ``dim != 2`` must not call these functions.  The D3Q19
equilibrium path lives in the CUDA kernel layer and is not mirrored on
the host at present.
"""

from __future__ import annotations

import numpy as np

# D2Q9 velocity vectors (must match descriptors.cuh / descriptors.py)
_CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.float32)
_CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.float32)
_W = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36],
              dtype=np.float32)


def compute_feq_d2q9(rho: float, ux: float, uy: float) -> np.ndarray:
    """Compute D2Q9 equilibrium distribution ``f_eq``.

    Args:
        rho: Macroscopic density.
        ux: x-velocity (lattice units).
        uy: y-velocity (lattice units).

    Returns:
        9-element float32 array ``[f0 .. f8]``.
    """
    u_sq = ux * ux + uy * uy
    feq = np.empty(9, dtype=np.float32)
    for i in range(9):
        cu = _CX[i] * ux + _CY[i] * uy
        feq[i] = rho * _W[i] * (
            1.0 + 3.0 * cu + 4.5 * cu * cu - 1.5 * u_sq)
    return feq


def compute_macro_from_ddf(f: np.ndarray) -> tuple[float, float, float]:
    """Compute rho, ux, uy from a 9-element D2Q9 DDF array.

    Args:
        f: 9-element float32 array of distribution functions.

    Returns:
        ``(rho, ux, uy)`` tuple.
    """
    rho = float(np.sum(f))
    rho_safe = rho if abs(rho) > 1e-12 else 1.0
    ux = float(np.sum(f * _CX)) / rho_safe
    uy = float(np.sum(f * _CY)) / rho_safe
    return rho, ux, uy
