# 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
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()

        # 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)

        # Snapshot
        self._ddf_snap: 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 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

    # -- DDF download with precision decode ----------------------------------
    def download_ddf(self):
        """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."""
        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]

    def _read_lattice_weights(self):
        """Read d_w[] from CUDA __constant__ memory (single source of truth in descriptors.cuh)."""
        ptr, size = self.module.get_global("d_w")
        w = np.empty(self.nq, dtype=np.float32)
        cuda.memcpy_dtoh(w, ptr)
        return 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):
        if self.store_bytes == 2:
            if self.store_dtype != np.float16:
                raise RuntimeError(
                    "Invalid 2-byte DDF storage dtype; expected float16 for FP16S."
                )
            buf = (self.ddf * 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, self.ddf)
            cuda.memcpy_htod(self.temp_gpu, self.ddf)

    # 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 cut-link and sensor compact lists to GPU."""
        self.curved.upload(stream=stream)
        self.sensors.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."""
        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)
            return cx, cy, cz
        return cx, cy

    # -- Macroscopic field extraction ----------------------------------------
    def get_macroscopic(self):
        """Download DDF and compute rho, ux, uy [, uz] on host."""
        self.download_ddf()
        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}")

    # -- Snapshots -----------------------------------------------------------
    def snapshot(self):
        self.download_ddf()
        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()
