# CelerisLab/lbm/curved_links.py
"""Structure-of-arrays helpers for curved Bouzidi links and sensor probe lists.

Responsibilities:
    Hold host columns and lazy GPU mirrors for compact boundary kernels.

Design::
    Device columns are reused by capacity to avoid reallocating every upload.
    ``count`` tracks active rows while ``capacity`` tracks allocated rows.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import pycuda.driver as cuda


def _free_alloc(ptr: Optional[cuda.DeviceAllocation]) -> None:
    if ptr is not None:
        ptr.free()


@dataclass
class CurvedLinkSoA:
    """Cut-link columns on host with matching ``*_gpu`` device buffers."""

    fluid_idx: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.uint32))
    dir: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.uint8))
    q: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.float32))
    body_id: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.int32))
    rx: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.float32))
    ry: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.float32))
    rz: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.float32))
    fallback_class: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.uint8))

    fluid_idx_gpu: Optional[cuda.DeviceAllocation] = None
    dir_gpu: Optional[cuda.DeviceAllocation] = None
    q_gpu: Optional[cuda.DeviceAllocation] = None
    body_id_gpu: Optional[cuda.DeviceAllocation] = None
    rx_gpu: Optional[cuda.DeviceAllocation] = None
    ry_gpu: Optional[cuda.DeviceAllocation] = None
    rz_gpu: Optional[cuda.DeviceAllocation] = None
    fallback_class_gpu: Optional[cuda.DeviceAllocation] = None

    count: int = 0
    capacity: int = 0

    def assign_host(
        self,
        fluid_idx: np.ndarray,
        dir: np.ndarray,
        q: np.ndarray,
        body_id: np.ndarray,
        rx: np.ndarray,
        ry: np.ndarray,
        rz: np.ndarray,
        fallback_class: np.ndarray,
    ) -> None:
        """Store and validate host SoA columns for curved-link kernel inputs."""
        expected_dtypes = {
            "fluid_idx": np.uint32,
            "dir": np.uint8,
            "q": np.float32,
            "body_id": np.int32,
            "rx": np.float32,
            "ry": np.float32,
            "rz": np.float32,
            "fallback_class": np.uint8,
        }
        raw_columns = {
            "fluid_idx": fluid_idx,
            "dir": dir,
            "q": q,
            "body_id": body_id,
            "rx": rx,
            "ry": ry,
            "rz": rz,
            "fallback_class": fallback_class,
        }

        columns = {}
        for name, value in raw_columns.items():
            arr = np.asarray(value)
            if not np.can_cast(arr.dtype, expected_dtypes[name], casting="same_kind"):
                raise TypeError(
                    f"CurvedLinkSoA column '{name}' dtype {arr.dtype} "
                    f"is incompatible with {np.dtype(expected_dtypes[name])}"
                )
            columns[name] = arr.astype(expected_dtypes[name], copy=False)

        n = len(columns["fluid_idx"])
        for name, col in columns.items():
            if len(col) != n:
                raise ValueError(
                    f"CurvedLinkSoA column '{name}' has length {len(col)} != {n}"
                )
            if not col.flags["C_CONTIGUOUS"]:
                columns[name] = np.ascontiguousarray(col)

        self.fluid_idx = columns["fluid_idx"]
        self.dir = columns["dir"]
        self.q = columns["q"]
        self.body_id = columns["body_id"]
        self.rx = columns["rx"]
        self.ry = columns["ry"]
        self.rz = columns["rz"]
        self.fallback_class = columns["fallback_class"]

        if self.q.size > 0:
            if np.any((self.q <= 0.0) | (self.q > 1.0)):
                raise ValueError("CurvedLinkSoA column 'q' must satisfy 0 < q <= 1")
        if self.fallback_class.size > 0:
            if not np.all(np.isin(self.fallback_class, (0, 1))):
                raise ValueError(
                    "CurvedLinkSoA column 'fallback_class' must use classes {0, 1}"
                )

    def upload(self, stream: Optional[cuda.Stream] = None) -> None:
        """Upload host columns to GPU with capacity-based allocation reuse.

        Args:
            stream: If set, use async host-to-device copies on this stream.
        """
        n = int(len(self.fluid_idx))
        self.count = n
        if n == 0:
            return

        if n > self.capacity or self.fluid_idx_gpu is None:
            self.free_gpu_columns()
            self.fluid_idx_gpu = cuda.mem_alloc(int(n * self.fluid_idx.dtype.itemsize))
            self.dir_gpu = cuda.mem_alloc(int(n * self.dir.dtype.itemsize))
            self.q_gpu = cuda.mem_alloc(int(n * self.q.dtype.itemsize))
            self.body_id_gpu = cuda.mem_alloc(int(n * self.body_id.dtype.itemsize))
            self.rx_gpu = cuda.mem_alloc(int(n * self.rx.dtype.itemsize))
            self.ry_gpu = cuda.mem_alloc(int(n * self.ry.dtype.itemsize))
            self.rz_gpu = cuda.mem_alloc(int(n * self.rz.dtype.itemsize))
            self.fallback_class_gpu = cuda.mem_alloc(
                int(n * self.fallback_class.dtype.itemsize)
            )
            self.capacity = n

        assert self.fluid_idx_gpu is not None
        assert self.dir_gpu is not None
        assert self.q_gpu is not None
        assert self.body_id_gpu is not None
        assert self.rx_gpu is not None
        assert self.ry_gpu is not None
        assert self.rz_gpu is not None
        assert self.fallback_class_gpu is not None

        if stream is not None:
            cuda.memcpy_htod_async(self.fluid_idx_gpu, self.fluid_idx, stream)
            cuda.memcpy_htod_async(self.dir_gpu, self.dir, stream)
            cuda.memcpy_htod_async(self.q_gpu, self.q, stream)
            cuda.memcpy_htod_async(self.body_id_gpu, self.body_id, stream)
            cuda.memcpy_htod_async(self.rx_gpu, self.rx, stream)
            cuda.memcpy_htod_async(self.ry_gpu, self.ry, stream)
            cuda.memcpy_htod_async(self.rz_gpu, self.rz, stream)
            cuda.memcpy_htod_async(
                self.fallback_class_gpu, self.fallback_class, stream)
        else:
            cuda.memcpy_htod(self.fluid_idx_gpu, self.fluid_idx)
            cuda.memcpy_htod(self.dir_gpu, self.dir)
            cuda.memcpy_htod(self.q_gpu, self.q)
            cuda.memcpy_htod(self.body_id_gpu, self.body_id)
            cuda.memcpy_htod(self.rx_gpu, self.rx)
            cuda.memcpy_htod(self.ry_gpu, self.ry)
            cuda.memcpy_htod(self.rz_gpu, self.rz)
            cuda.memcpy_htod(self.fallback_class_gpu, self.fallback_class)

    def free_gpu_columns(self) -> None:
        """Release device columns only (host arrays unchanged)."""
        _free_alloc(self.fluid_idx_gpu)
        _free_alloc(self.dir_gpu)
        _free_alloc(self.q_gpu)
        _free_alloc(self.body_id_gpu)
        _free_alloc(self.rx_gpu)
        _free_alloc(self.ry_gpu)
        _free_alloc(self.rz_gpu)
        _free_alloc(self.fallback_class_gpu)
        self.fluid_idx_gpu = None
        self.dir_gpu = None
        self.q_gpu = None
        self.body_id_gpu = None
        self.rx_gpu = None
        self.ry_gpu = None
        self.rz_gpu = None
        self.fallback_class_gpu = None
        self.capacity = 0

    def free(self) -> None:
        """Free GPU storage and reset ``count``."""
        self.free_gpu_columns()
        self.count = 0


@dataclass
class SensorSoA:
    """Sensor cell list with GPU mirrors for ``SensorKernel``."""

    cells: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.uint32))
    obj_id: np.ndarray = field(
        default_factory=lambda: np.zeros(0, dtype=np.int32))

    cells_gpu: Optional[cuda.DeviceAllocation] = None
    obj_id_gpu: Optional[cuda.DeviceAllocation] = None

    count: int = 0

    def assign_host(self, cells: np.ndarray, obj_id: np.ndarray) -> None:
        """Store host sensor columns."""
        self.cells = cells
        self.obj_id = obj_id

    def upload(self, stream: Optional[cuda.Stream] = None) -> None:
        """Reallocate GPU columns and upload.

        Args:
            stream: If set, use async host-to-device copies on this stream.
        """
        self.free_gpu_columns()
        n = int(len(self.cells))
        self.count = n
        if n == 0:
            return

        self.cells_gpu = cuda.mem_alloc(int(self.cells.nbytes))
        self.obj_id_gpu = cuda.mem_alloc(int(self.obj_id.nbytes))
        if stream is not None:
            cuda.memcpy_htod_async(self.cells_gpu, self.cells, stream)
            cuda.memcpy_htod_async(self.obj_id_gpu, self.obj_id, stream)
        else:
            cuda.memcpy_htod(self.cells_gpu, self.cells)
            cuda.memcpy_htod(self.obj_id_gpu, self.obj_id)

    def free_gpu_columns(self) -> None:
        _free_alloc(self.cells_gpu)
        _free_alloc(self.obj_id_gpu)
        self.cells_gpu = None
        self.obj_id_gpu = None

    def free(self) -> None:
        """Free GPU storage and reset ``count``."""
        self.free_gpu_columns()
        self.count = 0
