"""BodySyncPlan construction — build_flags_for, build_compact_lists_for, build_next_objects, build_sync_plan, commit_pending.

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.body.sync_plan import BodySyncPlan
from CelerisLab.lbm.descriptors import FLUID, SOLID, OBSTACLE, BC_CURVED


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


class _StubCfg:
    def __init__(self, dim=2):
        self.dim = dim
        self.is_d3q19 = (dim == 3)


class _StubField:
    """Minimal LBMField stub for build_sync_plan testing."""

    def __init__(self, nx: int, ny: int):
        self.nx = nx
        self.ny = ny
        # Build a simple channel flag array (fluid everywhere except top/bottom walls).
        n = nx * ny
        self.flag = np.ones(n, dtype=np.uint16) * FLUID
        self.flag[:nx] = SOLID | 0x0010  # bottom wall
        self.flag[(ny - 1) * nx:ny * nx] = SOLID | 0x0010  # top wall
        # Save a clean copy for build_channel_flags
        self._channel_flags = self.flag.copy()

    def build_channel_flags(self) -> np.ndarray:
        """Return a clean channel base (no object overlays)."""
        return self._channel_flags.copy()


NX, NY = 64, 32


class TestBuildFlagsFor(unittest.TestCase):

    def setUp(self):
        self.cfg = _StubCfg()
        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
        self.field = _StubField(NX, NY)

    def test_empty_objects_returns_base(self):
        base = self.field.build_channel_flags()
        result = ObjectManager.build_flags_for(
            [], base, nx=NX, ny=NY, nz=1)
        np.testing.assert_array_equal(result, base)

    def test_one_circle_produces_solid_region(self):
        base = self.field.build_channel_flags()
        obj = _make_circle_obj(32, 16, 5)
        result = ObjectManager.build_flags_for(
            [obj], base, nx=NX, ny=NY, nz=1)
        # Center cell should be solid with OBSTACLE and BC_CURVED bits
        center_idx = 32 + 16 * NX
        self.assertTrue(result[center_idx] & SOLID)
        self.assertTrue(result[center_idx] & OBSTACLE)

    def test_instance_method_delegates(self):
        base = self.field.build_channel_flags()
        obj = _make_circle_obj(32, 16, 5)
        self.mgr.add(obj)
        result = self.mgr.build_flags(base)
        center_idx = 32 + 16 * NX
        self.assertTrue(result[center_idx] & OBSTACLE)


class TestBuildCompactListsFor(unittest.TestCase):

    def setUp(self):
        self.cfg = _StubCfg()
        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)

    def test_circle_produces_curved_links(self):
        obj = _make_circle_obj(32, 16, 5)
        obj.obj_id = 0
        result = self.mgr.build_compact_lists_for([obj])
        cl_fluid_idx = result[0]
        self.assertGreater(len(cl_fluid_idx), 0)

    def test_sensor_produces_sensor_cells(self):
        obj = _make_circle_obj(32, 16, 5, is_sensor=True)
        obj.obj_id = 0
        result = self.mgr.build_compact_lists_for([obj])
        sensor_cells = result[8]
        self.assertGreater(len(sensor_cells), 0)

    def test_instance_method_delegates(self):
        obj = _make_circle_obj(32, 16, 5)
        self.mgr.add(obj)
        result = self.mgr.build_compact_lists()
        cl_fluid_idx = result[0]
        self.assertGreater(len(cl_fluid_idx), 0)


class TestBuildNextObjects(unittest.TestCase):

    def setUp(self):
        self.cfg = _StubCfg()
        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)

    def test_no_pending_returns_formal_objects(self):
        obj = _make_circle_obj(32, 16, 5)
        self.mgr.add(obj)
        result = self.mgr.build_next_objects()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].obj_id, 0)

    def test_removal_excludes_object(self):
        obj0 = _make_circle_obj(20, 16, 3)
        obj1 = _make_circle_obj(40, 16, 3)
        id0 = self.mgr.add(obj0)
        self.mgr.add(obj1)
        self.mgr.stage_remove(id0)
        result = self.mgr.build_next_objects()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].obj_id, 0)
        # The remaining object should be the second one (center at 40)
        self.assertAlmostEqual(result[0].center[0], 40.0)

    def test_add_appends_new_object(self):
        obj0 = _make_circle_obj(20, 16, 3)
        self.mgr.add(obj0)
        new_obj = _make_circle_obj(40, 16, 3)
        self.mgr.stage_add(new_obj)
        result = self.mgr.build_next_objects()
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0].obj_id, 0)
        self.assertEqual(result[1].obj_id, 1)

    def test_ids_are_consecutive(self):
        obj0 = _make_circle_obj(20, 16, 3)
        obj1 = _make_circle_obj(30, 16, 3)
        obj2 = _make_circle_obj(40, 16, 3)
        id0 = self.mgr.add(obj0)
        self.mgr.add(obj1)
        self.mgr.add(obj2)
        # Remove middle object
        self.mgr.stage_remove(id0 + 1)
        result = self.mgr.build_next_objects()
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0].obj_id, 0)
        self.assertEqual(result[1].obj_id, 1)

    def test_formal_registry_unchanged(self):
        obj = _make_circle_obj(32, 16, 5)
        id0 = self.mgr.add(obj)
        self.mgr.stage_remove(id0)
        self.mgr.build_next_objects()
        # Formal registry should be untouched
        self.assertEqual(self.mgr.count, 1)


class TestBuildSyncPlan(unittest.TestCase):

    def setUp(self):
        self.cfg = _StubCfg()
        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)
        self.field = _StubField(NX, NY)

    def test_add_body_produces_added_solid_mask(self):
        new_obj = _make_circle_obj(32, 16, 5)
        self.mgr.stage_add(new_obj)
        plan = self.mgr.build_sync_plan(self.field)
        self.assertIsInstance(plan, BodySyncPlan)
        self.assertEqual(plan.next_count, 1)
        # Center should be in added_solid_mask (was fluid, becomes solid)
        center_idx = 32 + 16 * NX
        self.assertTrue(plan.added_solid_mask[center_idx])

    def test_remove_body_produces_released_fluid_mask(self):
        obj = _make_circle_obj(32, 16, 5)
        id0 = self.mgr.add(obj)
        # Build current flags so the field "knows" about this body
        base = self.field.build_channel_flags()
        self.field.flag = self.mgr.build_flags(base)
        self.mgr.stage_remove(id0)
        plan = self.mgr.build_sync_plan(self.field)
        center_idx = 32 + 16 * NX
        self.assertTrue(plan.released_fluid_mask[center_idx])

    def test_no_change_masks_are_empty(self):
        plan = self.mgr.build_sync_plan(self.field)
        self.assertFalse(np.any(plan.added_solid_mask))
        self.assertFalse(np.any(plan.released_fluid_mask))


class TestCommitPending(unittest.TestCase):

    def setUp(self):
        self.cfg = _StubCfg()
        self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg)

    def test_commit_replaces_registry(self):
        obj0 = _make_circle_obj(20, 16, 3)
        obj1 = _make_circle_obj(40, 16, 3)
        id0 = self.mgr.add(obj0)
        self.mgr.add(obj1)
        self.mgr.stage_remove(id0)
        next_objs = self.mgr.build_next_objects()
        self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32))
        self.assertEqual(self.mgr.count, 1)
        self.assertEqual(self.mgr.get(0).obj_id, 0)
        self.assertFalse(self.mgr.has_pending_edit())

    def test_commit_clears_pending(self):
        self.mgr.stage_add(_make_circle_obj(32, 16, 5))
        next_objs = self.mgr.build_next_objects()
        self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32))
        self.assertFalse(self.mgr.has_pending_edit())


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