# CelerisLab/lbm/field.py
"""
LBMField — GPU memory for distribution functions, flags, and macroscopic data.

Owns:
  ddf_gpu / temp_gpu   – distribution-function double buffers
  flag_gpu             – cell-type flags
  curved / sensors     – SoA bundles for compact curved BC and sensor lists

Provides:
  upload_params()      – push LBMParams to __constant__ memory
  get_ddf() / set_ddf()
  get_macroscopic()    – rho, ux, uy [, uz]
  snapshot() / restore()
"""

import struct
from typing import Optional

import numpy as np
import pycuda.driver as cuda

from ..config import LBMConfig
from .curved_links import CurvedLinkSoA, SensorSoA, ForceRegionSoA, EsoPullCurvedBuffer
from .descriptors import FLUID, SOLID, BC_WALL, BC_INLET, BC_OUTLET

_SUPPORTED_STORE_PRECISIONS = ("FP32", "FP16S")


class LBMField:
    """Allocate and manage GPU arrays for one LBM domain."""

    def __init__(self, cfg: LBMConfig, module: cuda.Module):
        assert cfg.dim > 0 and cfg.nq > 0, (
            "cfg.validate() must be called before constructing LBMField")
        self.cfg = cfg
        self.module = module

        self.nx, self.ny, self.nz = cfg.nx, cfg.ny, cfg.nz
        self.nq = cfg.nq
        self.dim = cfg.dim
        self.n = self.nx * self.ny * self.nz
        self.dtype = np.float32  # compute precision always float32

        # Storage precision: current runtime path supports FP32 and FP16S only.
        _prec = getattr(cfg, 'store_precision', 'FP32')
        if _prec not in _SUPPORTED_STORE_PRECISIONS:
            supported = ", ".join(_SUPPORTED_STORE_PRECISIONS)
            raise RuntimeError(
                f"Unsupported store_precision '{_prec}' in LBMField. "
                f"Supported runtime precisions: {supported}."
            )
        self.store_bytes = 2 if _prec == 'FP16S' else 4
        self.store_dtype = np.float16 if _prec == 'FP16S' else np.float32

        # Host arrays (always float32 for convenience)
        self.ddf = np.zeros(self.n * self.nq, dtype=self.dtype)
        self.flag = self.build_channel_flags()

        # Compact lists (filled by ObjectManager.sync_to_gpu)
        self.curved = CurvedLinkSoA()
        self.sensors = SensorSoA()
        self.force_regions = ForceRegionSoA()
        self.esopull_curved = EsoPullCurvedBuffer()

        # GPU allocations – sized by storage precision
        _ddf_bytes = self.n * self.nq * self.store_bytes
        self.ddf_gpu = cuda.mem_alloc(_ddf_bytes)
        self.temp_gpu = cuda.mem_alloc(_ddf_bytes)
        self.flag_gpu = cuda.mem_alloc(self.flag.nbytes)

        # Macroscopic output buffers for EsoPull GPU diagnostics
        _macro_bytes = self.n * 4
        self.macro_rho_gpu = cuda.mem_alloc(_macro_bytes)
        self.macro_ux_gpu = cuda.mem_alloc(_macro_bytes)
        self.macro_uy_gpu = cuda.mem_alloc(_macro_bytes)
        self._macro_rho = np.zeros(self.n, dtype=np.float32)
        self._macro_ux = np.zeros(self.n, dtype=np.float32)
        self._macro_uy = np.zeros(self.n, dtype=np.float32)

        # Snapshot
        self._ddf_snap: np.ndarray | None = None
        self._host_ddf_step: int | None = None

        # Cached lattice constants mirrored from CUDA __constant__ memory.
        self._cached_w: np.ndarray | None = None
        self._cached_vectors_2d: tuple[np.ndarray, np.ndarray] | None = None
        self._cached_vectors_3d: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None

        # Host mirror of runtime parameters so partial updates preserve state.
        self._runtime_params = {
            "omega": float(cfg.omega),
            "fx": 0.0,
            "fy": 0.0,
            "fz": 0.0,
            "n_objects": 0,
        }

        # Upload d_params immediately
        self._upload_params()

    def invalidate_host_ddf_cache(self) -> None:
        """Mark host DDF mirror as stale after GPU-side updates."""
        self._host_ddf_step = None

    def invalidate_module_constant_cache(self) -> None:
        """Drop cached __constant__ mirrors after module replacement."""
        self._cached_w = None
        self._cached_vectors_2d = None
        self._cached_vectors_3d = None

    def build_channel_flags(self) -> np.ndarray:
        """Return a fresh host flag array for the base channel domain.

        The returned array contains only the static channel topology:
        inlet, outlet, y-walls, and interior fluid. Object overlays are applied
        later by ObjectManager on top of this clean base, which avoids stale
        obstacle bits when geometry is rebuilt or moved.
        """
        flag = np.ones(self.n, dtype=np.uint16) * FLUID

        if self.cfg.is_d2q9:
            flag[0:self.nx] = SOLID | BC_WALL
            flag[(self.ny - 1) * self.nx:self.ny * self.nx] = SOLID | BC_WALL
            flag[::self.nx] = SOLID | BC_INLET
            flag[self.nx - 1::self.nx] = SOLID | BC_OUTLET
            flag[0] = SOLID | BC_WALL
            flag[self.nx - 1] = SOLID | BC_WALL
            flag[(self.ny - 1) * self.nx] = SOLID | BC_WALL
            flag[self.ny * self.nx - 1] = SOLID | BC_WALL
            return flag

        flag = flag.reshape(self.nz, self.ny, self.nx)
        flag[:, 0, :] = SOLID | BC_WALL
        flag[:, self.ny - 1, :] = SOLID | BC_WALL
        flag[:, :, 0] = SOLID | BC_INLET
        flag[:, :, self.nx - 1] = SOLID | BC_OUTLET
        flag[:, 0, 0] = SOLID | BC_WALL
        flag[:, 0, self.nx - 1] = SOLID | BC_WALL
        flag[:, self.ny - 1, 0] = SOLID | BC_WALL
        flag[:, self.ny - 1, self.nx - 1] = SOLID | BC_WALL
        return flag.reshape(-1)

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

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

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

    # -- DDF download with precision decode ----------------------------------
    def download_ddf(self, *, step_id: int | None = None, force: bool = False):
        """Copy DDF from GPU to host self.ddf (float32), decoding FP16S if needed.
        If DDF shifting is enabled, un-shifts by adding back lattice weights."""
        # Fast path for repeated host queries within one solver step.
        # Callers pass step_id from LBMStepper; force=True bypasses cache
        # for snapshot-style semantics.
        if not force and step_id is not None and self._host_ddf_step == int(step_id):
            return

        if self.store_bytes == 2:
            buf = np.empty(self.n * self.nq, dtype=self.store_dtype)
            cuda.memcpy_dtoh(buf, self.ddf_gpu)
            if self.store_dtype == np.float16:
                # FP16S: stored as half(val * 32768) → decode by / 32768
                self.ddf[:] = buf.astype(np.float32) * 3.0517578125e-5
            else:
                self.ddf[:] = buf.astype(np.float32)
        else:
            cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)

        # Un-shift: kernel stores (f_i - w_i), restore to f_i for host analysis
        if getattr(self.cfg, 'ddf_shifting', False):
            w = self._read_lattice_weights()
            ddf_2d = self.ddf.reshape(self.nq, -1)
            for i in range(self.nq):
                ddf_2d[i] += w[i]

        # NOTE: EsoPull host DDF decode is NOT implemented here.
        # After EsoPullStep stores results, fi[k,i] is in EsoPull backing layout
        # (paired directions scattered across neighbor cells).  Reconstructing
        # physical DDF on the host requires per-cell neighbor tables.
        # Use get_macroscopic() which automatically selects the correct path
        # (host decode for double_buffer, GPU kernel for esopull).
        # See load_physical_node_esopull / store_physical_node_esopull.
        self._host_ddf_step = int(step_id) if step_id is not None else None

    def _read_lattice_weights(self):
        """Read d_w[] from CUDA __constant__ memory (single source of truth in descriptors.cuh)."""
        if self._cached_w is not None:
            return self._cached_w
        ptr, size = self.module.get_global("d_w")
        w = np.empty(self.nq, dtype=np.float32)
        cuda.memcpy_dtoh(w, ptr)
        self._cached_w = w
        return self._cached_w

    # -- d_params upload -----------------------------------------------------
    def _upload_params(self):
        """Pack LBMParams struct and upload to __constant__ d_params.

        Must match struct LBMParams in core/params.cuh layout:
            float omega;
            float fx, fy, fz;
            uint  n_objects;
        Total: 5 floats/uint words = 20 bytes.
        """
        p = self._runtime_params
        fmt = "=f fff I"
        data = struct.pack(
            fmt,
            p["omega"],
            p["fx"], p["fy"], p["fz"],
            p["n_objects"],
        )
        ptr, size = self.module.get_global("d_params")
        cuda.memcpy_htod(ptr, data)

    def update_params(self, **kwargs):
        """Re-upload d_params after changing runtime-adjustable values.

        Accepted keys: omega, fx, fy, fz, n_objects.
        """
        self._runtime_params.update({
            "omega": float(kwargs.get("omega", self._runtime_params["omega"])),
            "fx": float(kwargs.get("fx", self._runtime_params["fx"])),
            "fy": float(kwargs.get("fy", self._runtime_params["fy"])),
            "fz": float(kwargs.get("fz", self._runtime_params["fz"])),
            "n_objects": int(kwargs.get("n_objects", self._runtime_params["n_objects"])),
        })

        p = self._runtime_params
        fmt = "=f fff I"
        data = struct.pack(
            fmt,
            p["omega"],
            p["fx"], p["fy"], p["fz"],
            p["n_objects"],
        )
        ptr, _ = self.module.get_global("d_params")
        cuda.memcpy_htod(ptr, data)

    # -- Host ↔ GPU transfers ------------------------------------------------
    def upload_ddf(self):
        upload_data = self.ddf
        if getattr(self.cfg, 'ddf_shifting', False):
            # Host stores physical DDF; GPU stores shifted DDF (f - w).
            w = self._read_lattice_weights()
            upload_data = upload_data.copy()
            ddf_2d = upload_data.reshape(self.nq, -1)
            for i in range(self.nq):
                ddf_2d[i] -= w[i]

        if self.store_bytes == 2:
            if self.store_dtype != np.float16:
                raise RuntimeError(
                    "Invalid 2-byte DDF storage dtype; expected float16 for FP16S."
                )
            buf = (upload_data * 32768.0).astype(self.store_dtype)
            cuda.memcpy_htod(self.ddf_gpu, buf)
            cuda.memcpy_htod(self.temp_gpu, buf)
        else:
            cuda.memcpy_htod(self.ddf_gpu, upload_data)
            cuda.memcpy_htod(self.temp_gpu, upload_data)
        self.invalidate_host_ddf_cache()

    # The canonical download_ddf is defined above in the class body.

    def upload_flags(self):
        cuda.memcpy_htod(self.flag_gpu, self.flag)

    def upload_compact_lists(self, stream: Optional[cuda.Stream] = None) -> None:
        """Upload compact lists (curved, sensor, force_region) to GPU.

        Also allocates the EsoPull curved intermediate buffer when cut-links exist.
        The buffer is a GPU scratch space (no meaningful host data).
        """
        self.curved.upload(stream=stream)
        self.sensors.upload(stream=stream)
        self.force_regions.upload(stream=stream)
        # EsoPull curved scratch buffer: sized to n_curved, zero-filled.
        n = self.n_curved
        if n > 0:
            self.esopull_curved.assign_host(
                np.zeros(n, dtype=np.float32),
                np.zeros(n, dtype=np.uint32),
                np.zeros(n, dtype=np.uint8),
            )
        self.esopull_curved.upload(stream=stream)

    # -- Read lattice descriptors from CUDA module ---------------------------
    def _read_lattice_vectors(self):
        """Read d_cx, d_cy [, d_cz] from CUDA __constant__ memory."""
        if self.cfg.is_d2q9 and self._cached_vectors_2d is not None:
            return self._cached_vectors_2d
        if self.cfg.is_d3q19 and self._cached_vectors_3d is not None:
            return self._cached_vectors_3d

        nq = self.nq
        cx = np.empty(nq, dtype=np.int32)
        cy = np.empty(nq, dtype=np.int32)
        ptr_cx, _ = self.module.get_global("d_cx")
        ptr_cy, _ = self.module.get_global("d_cy")
        cuda.memcpy_dtoh(cx, ptr_cx)
        cuda.memcpy_dtoh(cy, ptr_cy)
        if self.cfg.is_d3q19:
            cz = np.empty(nq, dtype=np.int32)
            ptr_cz, _ = self.module.get_global("d_cz")
            cuda.memcpy_dtoh(cz, ptr_cz)
            self._cached_vectors_3d = (cx, cy, cz)
            return self._cached_vectors_3d
        self._cached_vectors_2d = (cx, cy)
        return self._cached_vectors_2d

    # -- Macroscopic field extraction ----------------------------------------
    def get_macroscopic(self, *, step_id: int | None = None):
        """Download DDF and compute rho, ux, uy [, uz].

        For double_buffer streaming: host-side decode from raw DDF.
        For esopull streaming: GPU-side MacroscopicEsoPullKernel.

        .. note::
            In EsoPull mode, ``get_ddf()`` returns backing-layout raw storage,
            NOT physical DDF.  Only ``get_macroscopic()`` is semantically correct.
            Verified: Kan99b K2 (200k steps, MRT) -- bit-identical to
            double-buffer for all metrics (St, mean_CL, mean_CD, amp_CL, amp_CD).
        """
        if getattr(self.cfg, 'streaming', 'double_buffer') == 'esopull':
            return self.get_macroscopic_esopull(step_id=step_id)

        # Double-buffer host path (existing)
        self.download_ddf(step_id=step_id)
        nq = self.nq

        if self.cfg.is_d2q9:
            f = self.ddf.reshape(nq, self.ny, self.nx)
            rho = np.sum(f, axis=0)
            cx, cy = self._read_lattice_vectors()
            rho_safe = np.where(np.abs(rho) > 1e-12, rho, 1.0)
            ux = sum(f[i] * cx[i] for i in range(nq)) / rho_safe
            uy = sum(f[i] * cy[i] for i in range(nq)) / rho_safe
            return {"rho": rho, "ux": ux, "uy": uy}

        if self.cfg.is_d3q19:
            f = self.ddf.reshape(nq, self.nz, self.ny, self.nx)
            rho = np.sum(f, axis=0)
            cx, cy, cz = self._read_lattice_vectors()
            rho_safe = np.where(np.abs(rho) > 1e-12, rho, 1.0)
            ux = sum(f[i] * cx[i] for i in range(nq)) / rho_safe
            uy = sum(f[i] * cy[i] for i in range(nq)) / rho_safe
            uz = sum(f[i] * cz[i] for i in range(nq)) / rho_safe
            return {"rho": rho, "ux": ux, "uy": uy, "uz": uz}

        raise ValueError(f"Unsupported lattice_model: {self.cfg.lattice_model}")

    def get_macroscopic_esopull(self, *, step_id: int | None = None) -> dict:
        """Compute macroscopic field from EsoPull backing layout on GPU.

        Launches MacroscopicEsoPullKernel to decode the EsoPull backing layout
        into physical rho/ux/uy, then downloads results to host.

        .. note::
            Only supports D2Q9.  D3Q19 raises NotImplementedError.
            This is the ONLY correct path for macroscopic diagnostics in EsoPull
            mode.  Host-side raw DDF download (get_ddf) is NOT physical DDF
            when streaming is esopull -- the backing layout scatters paired
            directions across neighbor cells.
            Verified: Kan99b K2 -- bit-identical to double-buffer.

        Args:
            step_id: Current step count used for EsoPull parity decoding.

        Returns:
            dict with keys ``"rho"``, ``"ux"``, ``"uy"`` as 2D arrays.
        """
        if not self.cfg.is_d2q9:
            raise NotImplementedError(
                "MacroscopicEsoPullKernel currently only supports D2Q9.")

        fn = self.module.get_function("MacroscopicEsoPullKernel")
        tpb = self.cfg.threads_per_block
        grid = ((self.nx + tpb - 1) // tpb, self.ny, 1)
        t_val = np.uint64(int(step_id) if step_id is not None else 0)

        fn(self.ddf_gpu,
           self.macro_rho_gpu, self.macro_ux_gpu, self.macro_uy_gpu,
           t_val,
           block=(tpb, 1, 1), grid=grid)

        cuda.memcpy_dtoh(self._macro_rho, self.macro_rho_gpu)
        cuda.memcpy_dtoh(self._macro_ux, self.macro_ux_gpu)
        cuda.memcpy_dtoh(self._macro_uy, self.macro_uy_gpu)

        return {
            "rho": self._macro_rho.reshape(self.ny, self.nx),
            "ux": self._macro_ux.reshape(self.ny, self.nx),
            "uy": self._macro_uy.reshape(self.ny, self.nx),
        }

    # -- DDF patch for runtime body sync (delegated to body/ddf_patch.py) -----
    # patch_ddf_for_body_sync and upload_patched_ddf are imported and called
    # directly from Simulation.sync_bodies().  They are no longer methods on
    # LBMField.

    def snapshot(self):
        self.download_ddf(force=True)
        self._ddf_snap = self.ddf.copy()

    def restore(self):
        if self._ddf_snap is None:
            raise RuntimeError("No snapshot to restore")
        self.ddf = self._ddf_snap.copy()
        self.upload_ddf()
