"""Checkpoint save -> load round trip for a scene with passive bodies. Requires GPU.

A fresh Simulation with the same scene, restored from the checkpoint must hold the same populations,
flags, body count and step count, and must stay bitwise identical to the original
after both continue for the same number of steps. Checkpoints with Jets are
refused (documented limitation) and are not covered here.
"""

import os
import tempfile
import unittest

import numpy as np

from ._cases import make_sim

# Three separated cylinders (pinball-like); cut-link receivers stay in bulk fluid.
CIRCLES = (((40, 32), 6), ((56, 22), 6), ((56, 42), 6))


def scene():
    sim = make_sim()
    sim.initialize()
    for center, radius in CIRCLES:
        sim.add_body("circle", center=center, radius=radius)
    sim.sync_bodies()
    return sim


class CheckpointBodiesGPU(unittest.TestCase):
    def test_round_trip_is_bitwise_and_continues_identically(self):
        sim = scene()
        sim.run(60)
        with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f:
            path = f.name
        restored = None
        try:
            saved = sim.save_checkpoint(path)
            # The loader requires the same body count and geometry: rebuild the scene.
            restored = scene()
            restored.load_checkpoint(saved)

            self.assertEqual(restored.bodies.count, len(CIRCLES))
            self.assertEqual(restored.stepper.step_count, sim.stepper.step_count)
            np.testing.assert_array_equal(restored.get_flags(), sim.get_flags())
            a, b = sim.get_ddf(), restored.get_ddf()
            self.assertEqual(a.dtype, b.dtype)
            self.assertEqual(a.tobytes(), b.tobytes())

            sim.run(40)
            restored.run(40)
            self.assertEqual(restored.stepper.step_count, sim.stepper.step_count)
            a, b = sim.get_ddf(), restored.get_ddf()
            self.assertTrue(np.isfinite(a).all())
            self.assertEqual(a.tobytes(), b.tobytes())
            for i in range(len(CIRCLES)):
                self.assertTrue(np.all(np.isfinite(restored.read_force(i))))
        finally:
            if restored is not None:
                restored.close()
            os.unlink(path)
            sim.close()


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