"""sync_bodies pipeline — recompile, esopull runtime body sync, step count preservation.

Requires GPU."""

import unittest
import tempfile
import os
import json

import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

from CelerisLab.simulation import Simulation, _scatter_to_field, _gather_from_field


# Use a small grid for fast compilation and test execution
NX, NY = 128, 64


class TestSyncBodiesSkeleton(unittest.TestCase):
    """Test sync_bodies() with real GPU -- double_buffer path."""

    def _make_sim(self) -> Simulation:
        """Create a Simulation with a small double_buffer D2Q9 grid."""
        return Simulation(device_id=0)

    def test_remove_body_sync_and_run(self):
        """Remove a body, sync, and continue running without crash."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Remove the body and sync
        sim.remove_body(0)
        sim.sync_bodies()

        # Should be able to continue running
        sim.run(50)

        # No body left -- count should be 0
        self.assertEqual(sim.bodies.count, 0)
        sim.close()

    def test_add_body_after_initialize(self):
        """Add a body after initialize, sync, and run."""
        sim = self._make_sim()
        sim.initialize()
        sim.run(50)

        # Add a body (returns -1 since it's staged)
        result_id = sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        self.assertEqual(result_id, -1)
        self.assertTrue(sim.bodies.has_pending_edit())

        # Sync commits the body
        sim.sync_bodies()
        self.assertFalse(sim.bodies.has_pending_edit())
        self.assertEqual(sim.bodies.count, 1)

        # Should be able to run with the new body
        sim.run(50)

        # Force readback should work
        force = sim.read_force(0)
        self.assertEqual(force.shape[0], 2)
        sim.close()

    def test_add_then_remove_body(self):
        """Add a body, then remove it, sync -- should result in zero bodies."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Add another body and remove the original
        sim.add_body("circle", center=(NX // 2, NY // 2), radius=6)
        sim.remove_body(0)
        sim.sync_bodies()

        # One body remaining (the newly added one, now id=0)
        self.assertEqual(sim.bodies.count, 1)
        sim.run(50)
        sim.close()

    def test_sync_preserves_step_count(self):
        """sync_bodies() should not reset the step counter."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(100)
        steps_before = sim.stepper.step_count

        sim.remove_body(0)
        sim.sync_bodies()

        steps_after = sim.stepper.step_count
        self.assertEqual(steps_after, steps_before)
        sim.close()

    def test_no_pending_edit_is_noop(self):
        """sync_bodies() with no pending edits should be a no-op."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # No edits -- sync should return immediately
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 1)
        sim.close()

    def test_run_discards_pending_without_sync(self):
        """Running without sync_bodies() should discard pending edits."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Stage a removal but don't sync
        sim.remove_body(0)
        self.assertTrue(sim.bodies.has_pending_edit())

        # run() auto-discards pending
        sim.run(50)
        self.assertFalse(sim.bodies.has_pending_edit())
        self.assertEqual(sim.bodies.count, 1)
        sim.close()


class TestSyncBodiesEsoPull(unittest.TestCase):
    """Test sync_bodies() with esopull streaming."""

    def _make_esopull_sim(self, nx=128, ny=64) -> Simulation:
        """Create a Simulation with a small esopull D2Q9 grid."""
        cfg = {
            "grid": {"lattice_model": "D2Q9", "nx": nx, "ny": ny, "nz": 1},
                "physics": {"data_type": "FP32", "viscosity": 0.05,
                            "velocity": 0.03, "rho": 1.0},
            "method": {
                "collision": "SRT", "streaming": "esopull",
                "store_precision": "FP32", "ddf_shifting": False,
                "les": {"enabled": False, "cs": 0.16, "closed_form": True},
                "trt": {"magic_param": 0.1875},
                "inlet": {"profile": "parabolic", "scheme": "zou_he_local"},
                "outlet": {"mode": "neq_extrap", "backflow_clamp": True,
                           "blend_alpha": 0.7},
                "y_wall_bc": "bounce_back",
                "omega_guard": {"min": 0.01, "max": 1.99},
            },
            "cuda": {"threads_per_block": 256, "compute_capability": "auto"},
        }
        tmpd = tempfile.mkdtemp(prefix="esopull_sync_")
        lbm_path = os.path.join(tmpd, "config_lbm.json")
        with open(lbm_path, "w") as f:
            json.dump(cfg, f)
        return Simulation(lbm_config_path=lbm_path)

    def test_esopull_add_body(self):
        """EsoPull: add body, sync, run, read finite force."""
        sim = self._make_esopull_sim()
        sim.initialize()
        sim.run(100)

        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 1)

        sim.run(100)
        force = sim.read_force(0)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"EsoPull sync: finite force expected, got {force}")
        sim.close()

    def test_esopull_remove_body(self):
        """EsoPull: add body → run → remove → sync → run."""
        sim = self._make_esopull_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(100)

        sim.remove_body(0)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 0)

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

    def test_esopull_add_remove_add(self):
        """EsoPull: add → remove → add cycle with finite checks."""
        sim = self._make_esopull_sim()
        sim.initialize()
        sim.run(100)

        # Add
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.sync_bodies()
        sim.run(100)

        # Remove
        sim.remove_body(0)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 0)
        sim.run(100)

        # Add again
        sim.add_body("circle", center=(NX // 2, NY // 2), radius=6)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 1)
        sim.run(100)
        force = sim.read_force(0)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"Finite force after esopull add-remove-add: {force}")
        sim.close()

    def test_esopull_roundtrip_fullgrid(self):
        """EsoPull: decode/encode roundtrip on full grid (no patch)."""
        sim = self._make_esopull_sim()
        sim.initialize()
        sim.run(100)

        nx, ny, nq = sim.field.nx, sim.field.ny, sim.field.nq
        all_cells = np.arange(nx * ny, dtype=np.uint32)
        step = sim.stepper.step_count
        stream = sim.stream

        s_gpu, s_host, idx_gpu = sim.field.esopull_sync_decode(
            all_cells, step, stream)
        _scatter_to_field(sim.field, s_host, all_cells)
        _gather_from_field(sim.field, s_host, all_cells)
        sim.field.esopull_sync_encode(
            s_gpu, s_host, idx_gpu, len(all_cells), step, stream)
        stream.synchronize()

        sim.run(100)
        macro = sim.get_macroscopic()
        self.assertTrue(np.all(np.isfinite(macro["ux"])),
                        "ux finite after full-grid roundtrip")
        sim.close()


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