"""Host contract of FixedPatchStepper(refinements=[]): one blocked level (mocked driver).

Per step: halo refresh -> static BFL (circles only) -> BlockStep with a null
kind, one stream wait per step(n) call, NaN spare after initialize. An upload,
launch or wait failure fails the instance and keeps its step count.
"""
from contextlib import nullcontext
from unittest.mock import Mock
import unittest

import numpy as np

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from tests.unit.test_refinement_three_level_multipatch import host_module


def circle(center=(15.5, 11.5), radius=4.3):
    return SimObject(0, geometry=CircleGeometry(*center, radius), center=center, radius=radius)


def backend(objects=(), failing_kernel=None, initialize=True):
    host, log = host_module(), []
    host.armed = False  # set True to make failing_kernel raise

    def launch(name):
        def run(*a, **k):
            log.append((name, a))
            if name == failing_kernel and host.armed:
                raise RuntimeError(name + " launch failed")
        return run

    module = Mock()
    module.get_global.return_value = (Mock(), 20)
    module.get_function.side_effect = lambda name: Mock(side_effect=launch(name))
    host.compiler.build_module.return_value = module
    host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size)
    host.cuda.memcpy_htod.side_effect = lambda pointer, array: log.append(("upload", (pointer, array)))
    run = host.FixedPatchStepper(LBMConfig(nx=32, ny=24, viscosity=.1),
                                 context=Mock(activate=nullcontext, sm_arch="sm_70"),
                                 refinements=[], level_blocks=[(12, 16)], level_objects=[objects])
    run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
    log.clear()
    if initialize:
        run.initialize(np.full(run._levels[0].shape, 1 / 9, np.float32))
    return run, log, host


def state():
    return np.full((9, 24, 32), 1 / 9, np.float32)


class SingleLevel(unittest.TestCase):
    def test_spare_starts_nan_and_current_holds_state(self):
        run, log, _ = backend()
        level = run._levels[0]
        uploads = {id(a[0]): a[1] for name, a in log if name == "upload"}
        self.assertTrue(np.isnan(uploads[id(level.spare)]).all())
        self.assertTrue(np.isfinite(level.layout.to_dense(uploads[id(level.current)])).all())

    def test_launch_order_per_step(self):
        for objects, order in (((), ["BlockHaloCopy", "BlockStep"]),
                               ((circle(),), ["BlockHaloCopy", "BlockCurvedBoundary", "BlockStep"])):
            with self.subTest(bodies=len(objects)):
                run, log, _ = backend(objects)
                log.clear()
                run.step(2)
                self.assertEqual([name for name, _ in log], order * 2 + ["sync"])
                step = [a for name, a in log if name == "BlockStep"][0]
                self.assertEqual(int(step[4]), 0)  # null kind: every core cell ACTIVE

    def test_circle_solids_are_not_state(self):
        run, _, _ = backend((circle(),))
        solid = run._level_boundaries[0].tables["solid"]
        self.assertTrue(solid.any())
        np.testing.assert_array_equal(run._state[0], ~solid)

    def test_initialize_upload_failure_fails_instance(self):
        run, _, host = backend(initialize=False)
        host.cuda.memcpy_htod.side_effect = RuntimeError("upload failed")
        with self.assertRaisesRegex(RuntimeError, "upload failed"):
            run.initialize(state())
        with self.assertRaisesRegex(RuntimeError, "failed"):
            run.step(1)

    def test_launch_failure_fails_instance_and_keeps_step_count(self):
        for kernel, objects in (("BlockHaloCopy", ()), ("BlockCurvedBoundary", (circle(),)),
                                ("BlockStep", ())):
            with self.subTest(kernel=kernel):
                run, log, host = backend(objects, failing_kernel=kernel)
                run.step(1)
                host.armed = True
                log.clear()
                with self.assertRaisesRegex(RuntimeError, kernel + " launch failed"):
                    run.step(2)
                self.assertEqual(run.step_count, 1)
                self.assertNotIn("sync", [name for name, _ in log])
                with self.assertRaisesRegex(RuntimeError, "failed"):
                    run.download()

    def test_final_wait_failure_fails_instance_and_keeps_step_count(self):
        run, _, _ = backend()
        run._stream.synchronize.side_effect = RuntimeError("device error")
        with self.assertRaisesRegex(RuntimeError, "device error"):
            run.step(3)
        self.assertEqual(run.step_count, 0)
        with self.assertRaisesRegex(RuntimeError, "failed"):
            run.step(1)


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