"""Full add/remove/checkpoint/load lifecycle — end-to-end body topology sync.

Requires GPU."""

import os
import unittest
import tempfile

import numpy as np
import pycuda.autoinit

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


class TestBodySyncE2E(unittest.TestCase):
    """Full end-to-end test of runtime body topology sync."""

    def test_full_lifecycle(self):
        """Create, init, run, add body, remove body, checkpoint, load."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        cx, cy = nx // 4, ny // 2

        # 1. Initialize and run
        sim.initialize()
        sim.run(100)
        self.assertGreater(sim.stepper.step_count, 0)

        # 2. Add a body and sync
        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)

        # 3. Run with body (short window — finite near-term)
        sim.run(50)
        force = sim.read_force(0)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"Finite force after add: {force}")

        steps_before_remove = sim.stepper.step_count

        # 4. Remove the body and sync
        sim.remove_body(0)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 0)
        self.assertTrue(sim.get_flags()[center_idx] & FLUID)

        # 5. Run after removal (no crash, step count advances)
        sim.run(50)
        self.assertEqual(sim.stepper.step_count, steps_before_remove + 50)

        # 6. Save checkpoint
        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f:
            ckpt_path = f.name
        try:
            saved_path = sim.save_checkpoint(ckpt_path)
            self.assertTrue(os.path.exists(saved_path))

            # 7. Load checkpoint in a new simulation
            sim2 = Simulation(device_id=0)
            sim2.initialize()
            sim2.run(1)

            sim2.load_checkpoint(saved_path)
            self.assertEqual(sim2.stepper.step_count, sim.stepper.step_count)
            self.assertEqual(sim2.bodies.count, 0)

            # 8. Continue running in restored sim
            sim2.run(50)
            # Verify step count advances (DDF may have NaN from pre-existing body)
            self.assertEqual(sim2.stepper.step_count,
                             sim.stepper.step_count + 50)
            sim2.close()
        finally:
            os.unlink(ckpt_path)

        sim.close()

    def test_add_remove_add_cycle(self):
        """Add → run → remove → run → add → run cycle with finite checks."""
        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(100)

        # Add
        sim.add_body("circle", center=(cx, cy), radius=8)
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 1)
        sim.run(50)

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

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


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