# CelerisLab/body/objects.py
"""
Lightweight flat object model for immersed / rigid bodies and sensors.

Design:
  - SimObject is a thin base with standard interface (state, control,
    get_flag_mask, get_curved_list).
  - Concrete types (Cylinder, Sensor, …) override geometry methods.
  - No deep inheritance tree.  Users can subclass SimObject directly
    for custom shapes — just implement get_flag_mask / get_curved_list.
"""

import math
import numpy as np
from dataclasses import dataclass
from typing import Tuple

# Lattice constants — canonical Python source is lbm/descriptors.py
from ..lbm.descriptors import (
    FLUID, SOLID, GAS, INTERFACE,
    BC_WALL, BC_INLET, BC_OUTLET, BC_CURVED, BC_PERIODIC, BC_MOVING,
    SENSOR_FLAG, OBSTACLE,
    D2Q9_EX, D2Q9_EY,
)

FALLBACK_CLASS_BOUZIDI = np.uint8(0)
FALLBACK_CLASS_HALFWAY_BOUNCEBACK = np.uint8(1)


# ---------------------------------------------------------------------------
# State / Control containers
# ---------------------------------------------------------------------------
@dataclass
class ObjectState:
    """Position, velocity, orientation (z used for 3D layouts)."""
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0
    theta: float = 0.0
    vx: float = 0.0
    vy: float = 0.0
    omega: float = 0.0


@dataclass
class ObjectControl:
    """Control input (force / velocity / displacement target)."""
    ax: float = 0.0
    ay: float = 0.0
    alpha: float = 0.0
    mode: int = 0  # 0=force, 1=velocity, 2=displacement


# ---------------------------------------------------------------------------
# Base class
# ---------------------------------------------------------------------------
class SimObject:
    """Base for all simulation objects."""

    obj_type: str = "generic"

    def __init__(self, obj_id: int, center: Tuple[float, ...],
                 radius: float = 0.0):
        self.obj_id = obj_id
        self.center = center
        self.radius = radius
        z0 = float(center[2]) if len(center) > 2 else 0.0
        self.state = ObjectState(x=float(center[0]), y=float(center[1]), z=z0)
        self.control = ObjectControl()

    def get_flag_mask(self, nx: int, ny: int, nz: int = 1) -> np.ndarray:
        """Return (n,) uint16 array with flag bits set for this object."""
        raise NotImplementedError




# ---------------------------------------------------------------------------
# Concrete types
# ---------------------------------------------------------------------------
class Cylinder(SimObject):
    obj_type = "cylinder"

    def get_flag_mask(self, nx: int, ny: int, nz: int = 1) -> np.ndarray:
        n = nx * ny * nz
        mask = np.zeros(n, dtype=np.uint16)
        xc, yc, zc = self.state.x, self.state.y, self.state.z
        r = self.radius
        r_sq = r * r
        x0 = max(1, int(xc - r) - 1)
        x1 = min(nx - 2, int(xc + r) + 1)
        y0 = max(1, int(yc - r) - 1)
        y1 = min(ny - 2, int(yc + r) + 1)
        z0 = max(1, int(zc - r) - 1) if nz > 1 else 0
        z1 = min(nz - 2, int(zc + r) + 1) if nz > 1 else 0
        if nz == 1:
            for x in range(x0, x1 + 1):
                for y in range(y0, y1 + 1):
                    if (x - xc) ** 2 + (y - yc) ** 2 < r_sq:
                        mask[x + y * nx] = SOLID | OBSTACLE | BC_CURVED
        else:
            for z in range(z0, z1 + 1):
                for y in range(y0, y1 + 1):
                    for x in range(x0, x1 + 1):
                        if (x - xc) ** 2 + (y - yc) ** 2 + (z - zc) ** 2 < r_sq:
                            k = x + y * nx + z * nx * ny
                            mask[k] = SOLID | OBSTACLE | BC_CURVED
        return mask

    def get_curved_list(self, nx: int, ny: int, nq: int,
                        nz: int = 1, ex=None, ey=None, ez=None,
                        domain_flags=None):
        """Return per-link curved boundary data (fluid-centric).

        Each record represents one cut link: a fluid node with a lattice
        direction that crosses the cylinder surface into solid.

        Args:
            nx, ny, nq: grid and lattice dimensions.
            nz: z-depth; use 1 for 2D layouts (linear index x + y·nx).
            ex, ey, ez: lattice vectors (length *nq*). *ez* required for a full
                D3Q19 fluid→solid link walk when ``nz > 1``.

        Returns:
            fluid_idx:  uint32  [n_links] fluid cell linear indices
            directions: uint8   [n_links] direction from fluid toward wall
            q_values:   float32 [n_links] fractional distance ∈ (0, 1]
            body_ids:   int32   [n_links] object id for each link
            rx_values:  float32 [n_links] x arm from center to wall hit point
            ry_values:  float32 [n_links] y arm from center to wall hit point
            rz_values:  float32 [n_links] z arm (zero for 2D)
            fallback:   uint8   [n_links]
                0=normal Bouzidi interpolation,
                1=half-way moving bounce-back fallback
        """
        from ..common.preprocess import (
            find_circle_intersection,
            find_sphere_ray_segment,
        )
        xc, yc, zc = self.state.x, self.state.y, self.state.z
        r = self.radius
        if ex is None or ey is None:
            ex, ey = D2Q9_EX, D2Q9_EY

        idx_list = []
        dir_list = []
        q_list = []
        rx_list = []
        ry_list = []
        rz_list = []
        fallback_list = []

        x0 = max(1, int(xc - r) - 2)
        x1 = min(nx - 2, int(xc + r) + 2)
        y0 = max(1, int(yc - r) - 2)
        y1 = min(ny - 2, int(yc + r) + 2)
        z0 = max(1, int(zc - r) - 2) if nz > 1 else 0
        z1 = min(nz - 2, int(zc + r) + 2) if nz > 1 else 0

        r_sq = r * r
        use_3d = ez is not None

        def _donor_is_fluid(xd: int, yd: int, zd: int = 0) -> bool:
            if domain_flags is None:
                return True
            if nz > 1:
                idx = xd + yd * nx + zd * nx * ny
            else:
                idx = xd + yd * nx
            return (int(domain_flags[idx]) & FLUID) != 0

        def _fallback_class(
            q_val: float, donor_inside: bool, donor_in_domain: bool, donor_is_fluid: bool
        ) -> np.uint8:
            """Classify donor legality for Bouzidi interpolation fallback."""
            if not (0.0 < q_val <= 1.0):
                raise ValueError(f"q must satisfy 0 < q <= 1, got {q_val}")
            if q_val >= 0.5:
                return FALLBACK_CLASS_BOUZIDI
            if (not donor_in_domain) or donor_inside or (not donor_is_fluid):
                return FALLBACK_CLASS_HALFWAY_BOUNCEBACK
            return FALLBACK_CLASS_BOUZIDI

        if use_3d:
            for z in range(z0, z1 + 1):
                for y in range(y0, y1 + 1):
                    for x in range(x0, x1 + 1):
                        if (x - xc) ** 2 + (y - yc) ** 2 + (z - zc) ** 2 < r_sq:
                            continue
                        for i in range(1, nq):
                            xn = x + ex[i]
                            yn = y + ey[i]
                            zn = z + ez[i]
                            if not (0 <= xn < nx and 0 <= yn < ny
                                    and 0 <= zn < nz):
                                continue
                            if (xn - xc) ** 2 + (yn - yc) ** 2 + (
                                    zn - zc) ** 2 >= r_sq:
                                continue
                            hit = find_sphere_ray_segment(
                                float(x), float(y), float(z),
                                float(xn), float(yn), float(zn),
                                xc, yc, zc, r,
                            )
                            if hit is None:
                                continue
                            dx = hit[0] - x
                            dy = hit[1] - y
                            dz = hit[2] - z
                            norm = math.sqrt(
                                ex[i] ** 2 + ey[i] ** 2 + ez[i] ** 2)
                            q = (math.sqrt(dx * dx + dy * dy + dz * dz) / norm
                                 if norm > 0 else 0.5)
                            q = max(0.01, min(q, 1.0))
                            xff = x - ex[i]
                            yff = y - ey[i]
                            zff = z - ez[i]
                            donor_in_domain = (0 <= xff < nx and 0 <= yff < ny and 0 <= zff < nz)
                            donor_inside = False
                            donor_is_fluid = False
                            if donor_in_domain:
                                donor_inside = (
                                    (xff - xc) ** 2 + (yff - yc) ** 2 + (zff - zc) ** 2
                                ) < r_sq
                                donor_is_fluid = _donor_is_fluid(xff, yff, zff)
                            idx_list.append(
                                np.uint32(x + y * nx + z * nx * ny))
                            dir_list.append(np.uint8(i))
                            q_list.append(np.float32(q))
                            rx_list.append(np.float32(hit[0] - xc))
                            ry_list.append(np.float32(hit[1] - yc))
                            rz_list.append(np.float32(hit[2] - zc))
                            fallback_list.append(
                                _fallback_class(q, donor_inside, donor_in_domain, donor_is_fluid))
        else:
            for x in range(x0, x1 + 1):
                for y in range(y0, y1 + 1):
                    if (x - xc) ** 2 + (y - yc) ** 2 < r_sq:
                        continue
                    for i in range(1, nq):
                        xn, yn = x + ex[i], y + ey[i]
                        if not (0 <= xn < nx and 0 <= yn < ny):
                            continue
                        if (xn - xc) ** 2 + (yn - yc) ** 2 >= r_sq:
                            continue
                        hit = find_circle_intersection(
                            float(x), float(y),
                            float(xn), float(yn),
                            xc, yc, r,
                        )
                        if hit is not None:
                            dx = hit[0] - x
                            dy = hit[1] - y
                            norm = math.sqrt(ex[i] ** 2 + ey[i] ** 2)
                            q = math.sqrt(dx * dx + dy * dy) / norm if norm > 0 else 0.5
                            q = max(0.01, min(q, 1.0))
                            xff = x - ex[i]
                            yff = y - ey[i]
                            donor_in_domain = (0 <= xff < nx and 0 <= yff < ny)
                            donor_inside = False
                            donor_is_fluid = False
                            if donor_in_domain:
                                donor_inside = ((xff - xc) ** 2 + (yff - yc) ** 2) < r_sq
                                donor_is_fluid = _donor_is_fluid(xff, yff)
                            idx_list.append(np.uint32(x + y * nx))
                            dir_list.append(np.uint8(i))
                            q_list.append(np.float32(q))
                            rx_list.append(np.float32(hit[0] - xc))
                            ry_list.append(np.float32(hit[1] - yc))
                            rz_list.append(np.float32(0.0))
                            fallback_list.append(
                                _fallback_class(q, donor_inside, donor_in_domain, donor_is_fluid))

        n = len(idx_list)
        if n > 0:
            return (np.array(idx_list, dtype=np.uint32),
                    np.array(dir_list, dtype=np.uint8),
                    np.array(q_list, dtype=np.float32),
                    np.full(n, self.obj_id, dtype=np.int32),
                    np.array(rx_list, dtype=np.float32),
                    np.array(ry_list, dtype=np.float32),
                    np.array(rz_list, dtype=np.float32),
                    np.array(fallback_list, dtype=np.uint8))
        return (np.zeros(0, dtype=np.uint32),
                np.zeros(0, dtype=np.uint8),
                np.zeros(0, dtype=np.float32),
                np.zeros(0, dtype=np.int32),
                np.zeros(0, dtype=np.float32),
                np.zeros(0, dtype=np.float32),
                np.zeros(0, dtype=np.float32),
                np.zeros(0, dtype=np.uint8))


class Sensor(SimObject):
    obj_type = "sensor"

    def get_flag_mask(self, nx: int, ny: int, nz: int = 1) -> np.ndarray:
        n = nx * ny * nz
        mask = np.zeros(n, dtype=np.uint16)
        xc, yc, zc = self.state.x, self.state.y, self.state.z
        r = self.radius
        r_sq = r * r
        x0 = max(1, int(xc - r) - 1)
        x1 = min(nx - 2, int(xc + r) + 1)
        y0 = max(1, int(yc - r) - 1)
        y1 = min(ny - 2, int(yc + r) + 1)
        z0 = max(1, int(zc - r) - 1) if nz > 1 else 0
        z1 = min(nz - 2, int(zc + r) + 1) if nz > 1 else 0
        if nz == 1:
            for x in range(x0, x1 + 1):
                for y in range(y0, y1 + 1):
                    if (x - xc) ** 2 + (y - yc) ** 2 < r_sq:
                        mask[x + y * nx] = FLUID | SENSOR_FLAG
        else:
            for z in range(z0, z1 + 1):
                for y in range(y0, y1 + 1):
                    for x in range(x0, x1 + 1):
                        if (x - xc) ** 2 + (y - yc) ** 2 + (z - zc) ** 2 < r_sq:
                            k = x + y * nx + z * nx * ny
                            mask[k] = FLUID | SENSOR_FLAG
        return mask

    def get_sensor_list(self, nx: int, ny: int, nz: int = 1):
        """Return compact sensor list.

        Returns:
            cells:   uint32 array of cell indices
            obj_ids: int32 array of object ids
        """
        cell_list = []
        xc, yc, zc = self.state.x, self.state.y, self.state.z
        r = self.radius
        r_sq = r * r
        x0 = max(1, int(xc - r) - 1)
        x1 = min(nx - 2, int(xc + r) + 1)
        y0 = max(1, int(yc - r) - 1)
        y1 = min(ny - 2, int(yc + r) + 1)
        z0 = max(1, int(zc - r) - 1) if nz > 1 else 0
        z1 = min(nz - 2, int(zc + r) + 1) if nz > 1 else 0
        if nz == 1:
            for x in range(x0, x1 + 1):
                for y in range(y0, y1 + 1):
                    if (x - xc) ** 2 + (y - yc) ** 2 < r_sq:
                        cell_list.append(np.uint32(x + y * nx))
        else:
            for z in range(z0, z1 + 1):
                for y in range(y0, y1 + 1):
                    for x in range(x0, x1 + 1):
                        if (x - xc) ** 2 + (y - yc) ** 2 + (z - zc) ** 2 < r_sq:
                            cell_list.append(
                                np.uint32(x + y * nx + z * nx * ny))
        if cell_list:
            cells = np.array(cell_list, dtype=np.uint32)
            obj_ids = np.full(len(cell_list), self.obj_id, dtype=np.int32)
            return cells, obj_ids
        return np.zeros(0, dtype=np.uint32), np.zeros(0, dtype=np.int32)