# CelerisLab/tests/unit/test_body_flags.py
"""Body type flag masks — OBSTACLE, SENSOR_FLAG, FRC_REGION bits for circle / sensor / force_region.

No GPU required."""

import unittest

import numpy as np

from CelerisLab.body.manager import ObjectManager
from CelerisLab.body.objects import SimObject
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.lbm.descriptors import (
    FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION,
)


def _make_obj(cx: float, cy: float, radius: float,
              is_sensor: bool = False,
              is_force_region: bool = False) -> SimObject:
    geom = CircleGeometry(cx, cy, radius)
    return SimObject(obj_id=-1, geometry=geom,
                     center=(cx, cy), radius=radius,
                     is_sensor=is_sensor,
                     is_force_region=is_force_region)


NX, NY = 64, 32


class TestBodyFlags(unittest.TestCase):
    """Verify flag mask bits for each body type."""

    def _obj_flag_mask(self, obj: SimObject) -> np.ndarray:
        return obj.get_flag_mask(NX, NY)

    def test_circle_has_obstacle_solid_curved(self):
        mask = self._obj_flag_mask(_make_obj(32, 16, 5))
        center = 32 + 16 * NX
        self.assertTrue(mask[center] & OBSTACLE,
                        "Circle should have OBSTACLE bit")
        self.assertTrue(mask[center] & SOLID,
                        "Circle should have SOLID bit")
        self.assertTrue(mask[center] & BC_CURVED,
                        "Circle should have BC_CURVED bit")
        self.assertFalse(mask[center] & FLUID,
                         "Circle interior should NOT be FLUID")

    def test_sensor_has_sensor_flag(self):
        mask = self._obj_flag_mask(_make_obj(32, 16, 5, is_sensor=True))
        center = 32 + 16 * NX
        self.assertTrue(mask[center] & SENSOR_FLAG,
                        "Sensor should have SENSOR_FLAG bit")
        self.assertTrue(mask[center] & FLUID,
                        "Sensor should be FLUID")

    def test_force_region_has_frc_region_flag(self):
        mask = self._obj_flag_mask(
            _make_obj(32, 16, 5, is_force_region=True))
        center = 32 + 16 * NX
        self.assertTrue(mask[center] & FRC_REGION,
                        "Force region should have FRC_REGION bit")
        self.assertTrue(mask[center] & FLUID,
                        "Force region should be FLUID")

    def test_circle_has_no_sensor_or_frc_flag(self):
        mask = self._obj_flag_mask(_make_obj(32, 16, 5))
        center = 32 + 16 * NX
        self.assertFalse(mask[center] & SENSOR_FLAG,
                         "Circle should NOT have SENSOR_FLAG")
        self.assertFalse(mask[center] & FRC_REGION,
                         "Circle should NOT have FRC_REGION")

    def test_force_region_has_no_obstacle(self):
        mask = self._obj_flag_mask(
            _make_obj(32, 16, 5, is_force_region=True))
        center = 32 + 16 * NX
        self.assertFalse(mask[center] & OBSTACLE,
                         "Force region should NOT have OBSTACLE bit")

    def test_all_body_types_nonzero_masks(self):
        for obj in [
            _make_obj(32, 16, 5),
            _make_obj(32, 16, 5, is_sensor=True),
            _make_obj(32, 16, 5, is_force_region=True),
        ]:
            mask = self._obj_flag_mask(obj)
            self.assertGreater(np.count_nonzero(mask), 0,
                               f"{obj.is_sensor=},{obj.is_force_region=}: "
                               "mask should have non-zero entries")


if __name__ == "__main__":
    unittest.main()