"""Unified action/obs flow — host-only set_body, auto transfer, stream API, DRL loop pattern.

Requires GPU."""

import unittest

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

from CelerisLab.simulation import Simulation


class TestUnifiedObs(unittest.TestCase):
    """Test unified action/obs flow."""

    def test_set_body_then_run_read_body(self):
        """set_body (host-only), run, read_body returns finite force."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # set_body should not trigger H2D (no error expected)
        sim.set_body(0, omega=0.001)

        # run will auto-upload action
        sim.run(50)
        data = sim.read_body(0)
        self.assertTrue(np.all(np.isfinite(data.force)),
                        f"Force finite: {data.force}")
        sim.close()

    def test_skip_transfer(self):
        """run(upload_act=False, sync_obs=False) should not crash."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
        sim.initialize()
        sim.run(50, upload_act=False, sync_obs=False)
        # After no-sync run, step count should still advance
        self.assertEqual(sim.stepper.step_count, 50)
        sim.close()

    def test_external_stream(self):
        """Providing an external CUDA stream should not crash."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
        sim.initialize()
        s = cuda.Stream()
        sim.run(50, stream=s)
        force = sim.read_force(0)
        self.assertTrue(np.all(np.isfinite(force)),
                        f"Force finite with external stream: {force}")
        sim.close()

    def test_read_body_before_run_returns_zeros(self):
        """read_body before any run() should return zero force (buffer is
        initialized to zero during sync_to_gpu)."""
        sim = Simulation(device_id=0)
        sim.add_body("circle", center=(128, 128), radius=8)
        sim.initialize()
        force = sim.read_force(0)
        np.testing.assert_array_equal(force, np.zeros(2, dtype=np.float32))
        sim.close()

    def test_drl_pattern(self):
        """DRL-style loop: run → read → set → run → read."""
        sim = Simulation(device_id=0)
        nx = sim.lbm_cfg.nx
        ny = sim.lbm_cfg.ny
        sim.add_body("circle", center=(nx // 4, ny // 2), radius=8)
        sim.initialize()
        for i in range(3):
            sim.run(50)
            data = sim.read_body(0)
            self.assertTrue(np.all(np.isfinite(data.force)))
            sim.set_body(0, omega=0.001 * i)
        self.assertEqual(sim.stepper.step_count, 150)
        sim.close()


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