# CelerisLab/body/ddf_patch.py
"""
Host-side DDF patching for runtime body topology sync. (D2Q9 only)

When bodies are added or removed at runtime, the host DDF must be
patched so that newly-solid cells contain equilibrium (not stale
fluid values) and newly-fluid cells initialise from their neighbours.

Called by ``Simulation.sync_bodies()`` after topology rebuild and
before GPU upload.

All equilibrium calculations use ``lbm/equilibrium.py`` which is D2Q9-only.
D3Q19 DDF patching is not yet implemented on the host side.
"""

from __future__ import annotations

from typing import Optional

import numpy as np
import pycuda.driver as cuda

from ..lbm.descriptors import FLUID
from ..lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf


def patch_ddf_for_body_sync(
    field,
    old_ddf: np.ndarray,
    old_flags: np.ndarray,
    new_flags: np.ndarray,
    added_solid_mask: np.ndarray,
    released_fluid_mask: np.ndarray,
    curved_fluid_indices: np.ndarray | None = None,
) -> None:
    """Patch host-side DDF after a body topology change.

    Operates on physical DDF (``field.ddf``).  After patching, the caller
    should call ``upload_patched_ddf()`` to push only the changed cells
    to both GPU buffers.

    Three patch types:

    * ``old_fluid -> new_solid``:  Write equilibrium at rest (rho=RHO, u=0)
      so stale fluid values do not linger inside new solid cells.

    * ``old_solid -> new_fluid``:  BFS inward fill from surrounding old-fluid
      cells.  Layer-0 cells average their old-fluid neighbours; deeper layers
      average the previously filled layer.  Fallback: equilibrium at inlet
      velocity.

    * **Boundary fluid cells** (new curved boundary neighbours):  If new
      solid cells appear, the fluid cells just outside them (on curved
      boundary links) are also set to equilibrium at their local velocity
      (or rest if NaN), preventing abrupt-boundary instability.

    Args:
        field: LBMField instance (``field.ddf`` is modified in-place).
        old_ddf: Physical DDF snapshot before topology change (n*NQ, float32).
        old_flags: Flag array before topology change.
        new_flags: Flag array after topology change.
        added_solid_mask: Bool array -- True where fluid became solid.
        released_fluid_mask: Bool array -- True where obstacle became fluid.
        curved_fluid_indices: Optional array of fluid cell indices on curved
            boundary links.  When *added_solid_mask* has any True entries,
            these cells are set to equilibrium at local velocity (or rest).
    """
    _patch_added_solid(field, added_solid_mask)
    _patch_boundary_fluid(field, curved_fluid_indices, added_solid_mask)
    _patch_released_fluid(field, old_ddf, old_flags, released_fluid_mask)


def upload_patched_ddf(
    field,
    added_solid_mask: np.ndarray,
    released_fluid_mask: np.ndarray,
    curved_fluid_indices: np.ndarray | None = None,
) -> None:
    """Upload only the DDF cells that changed during patching.

    Uses per-cell ``cuda.memcpy_htod`` for both ``ddf_gpu`` and
    ``temp_gpu``.  Leaves all other cells (especially boundary cells
    with NaN) untouched on the GPU, preserving their streaming-
    compatible values from before the sync.
    """
    changed_indices: set[int] = set()
    changed_indices.update(np.where(added_solid_mask)[0].tolist())
    changed_indices.update(np.where(released_fluid_mask)[0].tolist())
    if curved_fluid_indices is not None:
        changed_indices.update(curved_fluid_indices.tolist())
    if not changed_indices:
        field.invalidate_host_ddf_cache()
        field._host_ddf_step = None
        return

    nq = field.nq
    cell_bytes = nq * 4
    ddf_ptr = int(field.ddf_gpu)
    temp_ptr = int(field.temp_gpu)

    for idx in sorted(changed_indices):
        offset = idx * cell_bytes
        cell_data = field.ddf[idx * nq:(idx + 1) * nq].tobytes()
        cuda.memcpy_htod(ddf_ptr + offset, cell_data)
        cuda.memcpy_htod(temp_ptr + offset, cell_data)

    field.invalidate_host_ddf_cache()
    field._host_ddf_step = None


# -- Internal helpers ---------------------------------------------------------

def _patch_added_solid(field, added_solid_mask: np.ndarray) -> None:
    """Write resting equilibrium for cells that became solid."""
    if not np.any(added_solid_mask):
        return
    rho0 = field.cfg.rho
    feq = compute_feq_d2q9(rho0, 0.0, 0.0)
    indices = np.where(added_solid_mask)[0]
    nq = field.nq
    for idx in indices:
        field.ddf[idx * nq:(idx + 1) * nq] = feq


def _patch_boundary_fluid(
    field,
    curved_fluid_indices: np.ndarray | None,
    added_solid_mask: np.ndarray,
) -> None:
    """Patch fluid cells on new curved boundary links to local equilibrium."""
    if curved_fluid_indices is None or len(curved_fluid_indices) == 0:
        return
    if not np.any(added_solid_mask):
        return
    nq = field.nq
    feq_fallback = compute_feq_d2q9(field.cfg.rho, 0.0, 0.0)
    for idx in curved_fluid_indices:
        f = field.ddf[idx * nq:(idx + 1) * nq]
        if np.any(np.isnan(f)):
            field.ddf[idx * nq:(idx + 1) * nq] = feq_fallback
            continue
        rho, ux, uy = compute_macro_from_ddf(f)
        feq = compute_feq_d2q9(float(rho), float(ux), float(uy))
        field.ddf[idx * nq:(idx + 1) * nq] = feq


def _patch_released_fluid(
    field,
    old_ddf: np.ndarray,
    old_flags: np.ndarray,
    released_fluid_mask: np.ndarray,
) -> None:
    """BFS inward fill for cells that became fluid after body removal."""
    if not np.any(released_fluid_mask):
        return

    nx, ny = field.nx, field.ny
    nq = field.nq
    old_is_fluid = (old_flags & FLUID) != 0
    released = np.where(released_fluid_mask)[0]

    filled = np.zeros(field.n, dtype=bool)

    offsets = [
        (-1, -1), (-1, 0), (-1, 1),
        (0, -1),           (0, 1),
        (1, -1),  (1, 0),  (1, 1),
    ]

    # Precompute macroscopic from old_ddf for old-fluid cells
    # Skip cells with NaN (pre-existing body curvature may have
    # contaminated boundary-fluid DDF values).
    old_macro: dict[int, tuple[float, float, float]] = {}
    fluid_indices = np.where(old_is_fluid)[0]
    for idx in fluid_indices:
        f = old_ddf[idx * nq:(idx + 1) * nq]
        if np.any(np.isnan(f)):
            continue
        rho, ux, uy = compute_macro_from_ddf(f)
        old_macro[idx] = (rho, ux, uy)

    def _neighbours(idx: int):
        x = idx % nx
        y = idx // nx
        for dx, dy in offsets:
            xn, yn = x + dx, y + dy
            if 0 <= xn < nx and 0 <= yn < ny:
                yield xn + yn * nx

    prev_layer_filled: dict[int, tuple[float, float, float]] = {}
    remaining = set(released.tolist())

    for idx in released:
        nbrs = list(_neighbours(idx))
        fluid_nbrs = [n for n in nbrs if n in old_macro]
        if not fluid_nbrs:
            continue
        rho_avg = np.mean([old_macro[n][0] for n in fluid_nbrs])
        ux_avg = np.mean([old_macro[n][1] for n in fluid_nbrs])
        uy_avg = np.mean([old_macro[n][2] for n in fluid_nbrs])
        feq = compute_feq_d2q9(float(rho_avg), float(ux_avg), float(uy_avg))
        field.ddf[idx * nq:(idx + 1) * nq] = feq
        filled[idx] = True
        prev_layer_filled[idx] = (
            float(rho_avg), float(ux_avg), float(uy_avg))
        remaining.discard(idx)

    # BFS layers
    max_layers = max(nx, ny)
    for _ in range(max_layers):
        if not remaining:
            break
        current_layer: dict[int, tuple[float, float, float]] = {}
        for idx in list(remaining):
            nbrs = list(_neighbours(idx))
            filled_nbrs = [n for n in nbrs if n in prev_layer_filled]
            if not filled_nbrs:
                continue
            rho_avg = np.mean([prev_layer_filled[n][0] for n in filled_nbrs])
            ux_avg = np.mean([prev_layer_filled[n][1] for n in filled_nbrs])
            uy_avg = np.mean([prev_layer_filled[n][2] for n in filled_nbrs])
            feq = compute_feq_d2q9(
                float(rho_avg), float(ux_avg), float(uy_avg))
            field.ddf[idx * nq:(idx + 1) * nq] = feq
            filled[idx] = True
            current_layer[idx] = (
                float(rho_avg), float(ux_avg), float(uy_avg))
            remaining.discard(idx)
        prev_layer_filled = current_layer

    # Fallback for any remaining unfilled cells
    u_inlet = field.cfg.velocity
    for idx in remaining:
        y = idx // nx
        ux_fallback = (
            u_inlet * 4.0 * y * (ny - 1 - y) / ((ny - 1) ** 2)
            if ny > 2 else u_inlet)
        feq = compute_feq_d2q9(field.cfg.rho, float(ux_fallback), 0.0)
        field.ddf[idx * nq:(idx + 1) * nq] = feq
