"""DDF patch after add / remove body — finite forces, finite macroscopic field after moderate steps.

Requires GPU."""

import unittest

import numpy as np
import pycuda.autoinit

from CelerisLab.simulation import Simulation
from CelerisLab.lbm.descriptors import FLUID, OBSTACLE


class TestDDFPatchAddBody(unittest.TestCase):
    """Test adding a body (fluid -> solid DDF patch)."""

    def test_add_body_runs_stably(self):
        """After adding a body, the simulation runs with finite forces
        for a moderate number of steps."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        cx, cy = nx // 4, ny // 2

        sim.initialize()
        sim.run(200)

        sim.add_body("circle", center=(cx, cy), radius=8)
        sim.sync_bodies()

        self.assertEqual(sim.bodies.count, 1)
        center_idx = cx + cy * nx
        self.assertTrue(sim.get_flags()[center_idx] & OBSTACLE)

        sim.run(50)
        force = sim.read_force(0)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"Finite force after add body: {force}")
        sim.close()


class TestDDFPatchRemoveBody(unittest.TestCase):
    """Test removing a body (solid -> fluid DDF patch via BFS inward fill)."""

    def test_remove_body_released_region_is_fluid(self):
        """After removal, the former body center should be a fluid cell."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        cx, cy = nx // 4, ny // 2

        sim.add_body("circle", center=(cx, cy), radius=8)
        sim.initialize()
        sim.run(200)

        sim.remove_body(0)
        sim.sync_bodies()

        flags = sim.get_flags()
        center_idx = cx + cy * nx
        self.assertTrue(flags[center_idx] & FLUID)
        sim.close()

    def test_remove_body_finite_field(self):
        """After removal, the macroscopic field is finite."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        cx, cy = nx // 4, ny // 2

        sim.add_body("circle", center=(cx, cy), radius=8)
        sim.initialize()
        sim.run(50)

        sim.remove_body(0)
        sim.sync_bodies()

        flags = sim.get_flags()
        center_idx = cx + cy * nx
        self.assertTrue(flags[center_idx] & FLUID)

        sim.run(50)
        macro = sim.get_macroscopic()
        self.assertTrue(np.all(np.isfinite(macro["ux"])),
                        "Macroscopic ux should be finite after remove body")
        sim.close()


class TestDDFPatchAddAndRemove(unittest.TestCase):
    """Test combined add + remove in one sync."""

    def test_add_one_remove_another(self):
        """Add a body and remove a different one in the same sync."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        cx, cy = nx // 4, ny // 2

        sim.add_body("circle", center=(cx, cy), radius=8)
        sim.initialize()
        sim.run(50)

        sim.add_body("circle", center=(nx // 2, ny // 2), radius=6)
        sim.remove_body(0)
        sim.sync_bodies()

        self.assertEqual(sim.bodies.count, 1)

        sim.run(50)
        force = sim.read_force(0)
        self.assertEqual(force.shape[0], 2)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"Finite force after add+remove: {force}")
        sim.close()


class TestDDFPatchNoChange(unittest.TestCase):
    """Test that patch is a no-op when there are no geometry changes."""

    def test_no_mask_no_change(self):
        """When neither mask has any True entries, DDF should be unchanged."""
        sim = Simulation(device_id=0)
        sim.add_body("circle", center=(128, 128), radius=8)
        sim.initialize()
        sim.run(100)

        sim.field.download_ddf(force=True)
        ddf_before = sim.field.ddf.copy()

        from CelerisLab.body.ddf_patch import patch_ddf_for_body_sync as patch_fn

        n = sim.field.n
        added = np.zeros(n, dtype=bool)
        released = np.zeros(n, dtype=bool)
        patch_fn(
            sim.field,
            ddf_before, sim.field.flag.copy(), sim.field.flag.copy(),
            added, released)

        np.testing.assert_array_equal(sim.field.ddf, ddf_before)
        sim.close()


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