"""CPU geometry tests for lbm.hierarchy (no CUDA). Geometry construction only, not LBM numerics."""
import unittest

import numpy as np
from CelerisLab.config import LBMConfig
from CelerisLab.lbm import refinement as r
from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition

CFG = LBMConfig(nx=48, ny=48, viscosity=.02)
L1 = [(6, 6, 26, 26), (32, 30, 42, 42)]   # L0 cells; L1 ACTIVE x/y 12..52 and x 64..84, y 60..84
L2 = [(20, 20, 44, 44), (72, 68, 76, 76)]  # L1 cells; L2 ACTIVE 40..88 and x 144..152, y 136..152
L3 = [(56, 56, 72, 72)]                    # L2 cells
FOUR = [L1, L2, L3]
SMALL = LBMConfig(nx=16, ny=16, viscosity=.1)


class FourLevels(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.h = build_hierarchy(CFG, FOUR)

    def test_config_chain(self):
        self.assertEqual(self.h.depth, 4)
        for level, c in enumerate(self.h.configs):
            self.assertEqual((c.ny, c.nx), (48 << level, 48 << level))
            self.assertAlmostEqual(c.viscosity, .02 * 2 ** level, places=12)
            self.assertEqual(self.h.kinds[level].shape, (c.ny, c.nx))
        for level, p in enumerate(self.h.pairs):
            self.assertAlmostEqual(p.tau_coarse, 1 / self.h.configs[level].omega, places=12)
            self.assertAlmostEqual(p.tau_fine, 1 / self.h.configs[level + 1].omega, places=12)

    def test_partition_and_composition(self):
        h = self.h
        np.testing.assert_array_equal(leaf_partition(h), np.ones((48, 48)))
        for level in (1, 2):
            covered, raw = h.pairs[level].covered, h.pairs[level - 1].fine_kind
            self.assertTrue((h.kinds[level][covered] == r.OUTER_GHOST).all())
            np.testing.assert_array_equal(h.kinds[level][~covered], raw[~covered])
        np.testing.assert_array_equal(h.kinds[3], h.pairs[2].fine_kind)

    def test_depth_is_list_length(self):
        for depth in range(1, 5):
            h = build_hierarchy(CFG, FOUR[:depth - 1])
            self.assertEqual(h.depth, depth)
            np.testing.assert_array_equal(leaf_partition(h), np.ones((48, 48)))
            for level in range(depth - 1):
                np.testing.assert_array_equal(h.kinds[level], self.h.kinds[level])

    def test_illegal_coverage_reports_first_guard(self):
        cases = [([L1, [(10, 20, 30, 44)]], "inherited ACTIVE"),
                 ([L1, L2[:1] + [(54, 54, 58, 58)]], "inherited ACTIVE"),
                 ([L1, L2, [(40, 40, 60, 60)]], "donors/receivers"),
                 ([L1, L2, [(41, 41, 60, 60)]], "donors/receivers")]
        for levels, reason in cases:
            with self.subTest(levels=levels), self.assertRaisesRegex(ValueError, reason):
                build_hierarchy(CFG, levels)
        build_hierarchy(CFG, [L1, L2, [(56, 56, 64, 64), (66, 56, 72, 64)]])  # gap 2 closes
        build_hierarchy(CFG, [L1, L2, [(42, 42, 60, 60)]])  # minimal legal margin


class LegacyForms(unittest.TestCase):
    def test_scalar_parent_and_child_match_existing_calls(self):
        local = r.FixedPatch2D.from_config(SMALL, (4, 4, 12, 12))
        child, kind = r.nest_patch(SMALL, local, (8, 8, 16, 16))
        h = build_hierarchy(SMALL, [(4, 4, 12, 12), (8, 8, 16, 16)])
        np.testing.assert_array_equal(h.kinds[0], np.where(local.covered, r.OUTER_GHOST, r.ACTIVE))
        np.testing.assert_array_equal(h.kinds[1], kind)
        np.testing.assert_array_equal(h.kinds[2], child.fine_kind)
        for name in ("ghost_fine", "ghost_coarse", "coalesce_coarse", "coalesce_fine"):
            np.testing.assert_array_equal(getattr(h.pairs[1], name), getattr(child, name))
        np.testing.assert_array_equal(leaf_partition(h), np.ones((16, 16)))

    def test_set_child_under_scalar_parent(self):
        local = r.FixedPatch2D.from_config(SMALL, (4, 4, 12, 12))
        child, kind = r.nest_patch_set(SMALL, local, [(8, 8, 16, 16)])
        h = build_hierarchy(SMALL, [(4, 4, 12, 12), [(8, 8, 16, 16)]])
        np.testing.assert_array_equal(h.kinds[1], kind)
        np.testing.assert_array_equal(h.pairs[1].coalesce_fine, child.coalesce_fine)
        np.testing.assert_array_equal(leaf_partition(h), np.ones((16, 16)))

    def test_scalar_child_under_set_parent_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "scalar parent"):
            build_hierarchy(SMALL, [[(4, 4, 12, 12)], (12, 12, 20, 20)])

    def test_wall_rows_and_single_level(self):
        h = build_hierarchy(SMALL, [(4, 4, 12, 12)], wall_aware=True)
        self.assertTrue((h.kinds[0][[0, -1]] == r.OUTER_GHOST).all())
        one = build_hierarchy(SMALL, [])
        self.assertEqual(one.depth, 1)
        self.assertTrue((one.kinds[0] == r.ACTIVE).all())
        walled = build_hierarchy(SMALL, [], wall_aware=True)  # regression: no refinement
        self.assertTrue((walled.kinds[0][[0, -1]] == r.OUTER_GHOST).all())
        self.assertTrue((walled.kinds[0][1:-1] == r.ACTIVE).all())


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