"""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.
"""
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=()):
    host, log = host_module(), []
    module = Mock()
    module.get_global.return_value = (Mock(), 20)
    module.get_function.side_effect = lambda name: Mock(
        side_effect=lambda *a, **k: log.append((name, a)))
    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()
    run.initialize(np.full(run._levels[0].shape, 1 / 9, np.float32))
    return run, log


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)


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