# CelerisLab/body/geometry/circle.py
"""
Circle geometry: cut-link enumeration, sensor cells, flag mask.

2D only.  For 3D, a separate SphereGeometry would be added.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple

import numpy as np

from .base import CutLink, Geometry, SensorCell
from ...lbm.descriptors import (
    FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION,
    D2Q9_EX, D2Q9_EY,
)


class CircleGeometry(Geometry):
    """Circular geometry for 2D immersed bodies and sensors."""

    def __init__(self, cx: float, cy: float, radius: float) -> None:
        self._cx = cx
        self._cy = cy
        self._radius = radius

    @property
    def center(self) -> Tuple[float, float]:
        return (self._cx, self._cy)

    # ------------------------------------------------------------------
    # Cut-link enumeration
    # ------------------------------------------------------------------
    def build_cut_links(
        self, nx: int, ny: int,
        domain_flags: Optional[np.ndarray] = None,
    ) -> List[CutLink]:
        """Enumerate cut links for this circle.

        Iterates the AABB bounding box around the circle, walks D2Q9
        directions from each fluid cell, records links that cross the
        circle boundary into solid.

        Args:
            nx, ny: Grid dimensions.
            domain_flags: Optional uint16 flags array for donor-cell
                fluid/solid classification. None = all donors assumed fluid.

        Returns:
            List of CutLink records (empty if zero links found).
        """
        from ...common.preprocess import find_circle_intersection

        cx, cy = self._cx, self._cy
        r = self._radius
        r_sq = r * r

        # AABB clip
        x0 = max(1, int(cx - r) - 2)
        x1 = min(nx - 2, int(cx + r) + 2)
        y0 = max(1, int(cy - r) - 2)
        y1 = min(ny - 2, int(cy + r) + 2)

        def _donor_is_fluid(xd: int, yd: int) -> bool:
            if domain_flags is None:
                return True
            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,
        ) -> int:
            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 0
            if (not donor_in_domain) or donor_inside or (not donor_is_fluid):
                return 1
            return 0

        links: List[CutLink] = []
        for x in range(x0, x1 + 1):
            for y in range(y0, y1 + 1):
                if (x - cx) ** 2 + (y - cy) ** 2 < r_sq:
                    continue
                for i in range(1, 9):  # D2Q9 non-zero directions
                    xn, yn = x + D2Q9_EX[i], y + D2Q9_EY[i]
                    if not (0 <= xn < nx and 0 <= yn < ny):
                        continue
                    if (xn - cx) ** 2 + (yn - cy) ** 2 >= r_sq:
                        continue
                    hit = find_circle_intersection(
                        float(x), float(y),
                        float(xn), float(yn),
                        cx, cy, r,
                    )
                    if hit is None:
                        continue
                    dx = hit[0] - x
                    dy = hit[1] - y
                    norm = math.sqrt(
                        float(D2Q9_EX[i]) ** 2 + float(D2Q9_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 - D2Q9_EX[i]
                    yff = y - D2Q9_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 - cx) ** 2 + (yff - cy) ** 2
                        ) < r_sq
                        donor_is_fluid = _donor_is_fluid(xff, yff)
                    links.append(CutLink(
                        fluid_idx=x + y * nx,
                        dir=i,
                        q=float(q),
                        body_id=0,  # assigned during packing
                        hit_x=float(hit[0]),
                        hit_y=float(hit[1]),
                        hit_z=0.0,
                        rx=float(hit[0] - cx),
                        ry=float(hit[1] - cy),
                        rz=0.0,
                        normal_x=0.0,
                        normal_y=0.0,
                        fallback=_fallback_class(
                            q, donor_inside, donor_in_domain, donor_is_fluid),
                        scheme_tag=0,
                        motion_tag=1,
                    ))

        return links

    # ------------------------------------------------------------------
    # Sensor cells
    # ------------------------------------------------------------------
    def build_sensor_cells(self, nx: int, ny: int) -> List[SensorCell]:
        """Enumerate sensor cells inside the circular footprint.

        Returns:
            List of SensorCell records.
        """
        cx, cy = self._cx, self._cy
        r = self._radius
        r_sq = r * r

        x0 = max(1, int(cx - r) - 1)
        x1 = min(nx - 2, int(cx + r) + 1)
        y0 = max(1, int(cy - r) - 1)
        y1 = min(ny - 2, int(cy + r) + 1)

        cells: List[SensorCell] = []
        for x in range(x0, x1 + 1):
            for y in range(y0, y1 + 1):
                if (x - cx) ** 2 + (y - cy) ** 2 < r_sq:
                    cells.append(SensorCell(idx=x + y * nx, body_id=0))
        return cells

    # ------------------------------------------------------------------
    # Flag mask
    # ------------------------------------------------------------------
    def build_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask with OBSTACLE|SOLID|BC_CURVED."""
        return self._build_flag_mask_inner(nx, ny, SOLID | OBSTACLE | BC_CURVED)

    def build_sensor_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask with FLUID|SENSOR_FLAG."""
        return self._build_flag_mask_inner(nx, ny, FLUID | SENSOR_FLAG)

    def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray:
        """Return (nx*ny,) uint16 flag mask with FLUID|FRC_REGION."""
        return self._build_flag_mask_inner(nx, ny, FLUID | FRC_REGION)

    def _build_flag_mask_inner(self, nx: int, ny: int, flag_val: int) -> np.ndarray:
        n = nx * ny
        mask = np.zeros(n, dtype=np.uint16)
        cx, cy = self._cx, self._cy
        r = self._radius
        r_sq = r * r
        x0 = max(1, int(cx - r) - 1)
        x1 = min(nx - 2, int(cx + r) + 1)
        y0 = max(1, int(cy - r) - 1)
        y1 = min(ny - 2, int(cy + r) + 1)
        for x in range(x0, x1 + 1):
            for y in range(y0, y1 + 1):
                if (x - cx) ** 2 + (y - cy) ** 2 < r_sq:
                    mask[x + y * nx] = np.uint16(flag_val)
        return mask
