"""CPU block layout contracts: ownership, halo pairs, packing and tagged pull."""
import sys
import unittest

import numpy as np

from CelerisLab.lbm.blocks import CORNER, FACE, NQ, BlockLayout2D
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY

CASES = (  # (ny, nx, by, bx, halo)
    (6, 8, 3, 4, 1),   # 2x2 split
    (6, 12, 3, 4, 1),  # 2x3 split
    (8, 8, 4, 4, 2),   # wider halo
    (4, 8, 4, 4, 1),   # one block along y
)
PERIODIC = ((True, True), (False, False), (True, False), (False, True))


def tags(ny, nx):
    """Unique, exactly representable value per dense population slot."""
    return np.arange(NQ * ny * nx, dtype=np.float64).reshape(NQ, ny, nx)


def dense_pull(a, periodic):
    """f[q](x) = a[q](x - c_q); NaN where a non-periodic exterior is crossed."""
    px, py = periodic
    _, ny, nx = a.shape
    gy, gx = np.mgrid[0:ny, 0:nx]
    out = np.empty_like(a)
    for q in range(NQ):
        sy, sx = gy - D2Q9_EY[q], gx - D2Q9_EX[q]
        valid = ((py | ((sy >= 0) & (sy < ny))) & (px | ((sx >= 0) & (sx < nx))))
        out[q] = np.where(valid, a[q, sy % ny, sx % nx], np.nan)
    return out


def blocked_pull(layout, blocked):
    """Exchange, then pull every core cell from its own block storage only."""
    layout.exchange(blocked)
    out = np.full_like(blocked, np.nan)
    (by, bx), h = layout.block, layout.halo
    for q in range(NQ):
        ey, ex = D2Q9_EY[q], D2Q9_EX[q]
        out[:, q, h:h + by, h:h + bx] = blocked[:, q, h - ey:h - ey + by, h - ex:h - ex + bx]
    return out


class BlockLayoutTests(unittest.TestCase):
    def test_cpu_import_boundary(self):
        self.assertFalse(any(name == "pycuda" or name.startswith("pycuda.") for name in sys.modules))

    def test_origins_and_unique_core_ownership(self):
        layout = BlockLayout2D.build(6, 8, 3, 4)
        np.testing.assert_array_equal(layout.origins, [[0, 0], [0, 4], [3, 0], [3, 4]])
        self.assertEqual(layout.blocked_shape, (4, NQ, 5, 6))
        for ny, nx, by, bx, h in CASES:
            layout = BlockLayout2D.build(ny, nx, by, bx, halo=h)
            ly, lx = layout.local_shape
            core = [b * ly * lx + (y + h) * lx + x + h
                    for b in range(layout.nblocks) for y in range(by) for x in range(bx)]
            self.assertEqual(sorted(layout.owner.ravel().tolist()), sorted(core))
            for y in range(ny):
                for x in range(nx):
                    b = (y // by) * (nx // bx) + x // bx
                    self.assertEqual(layout.block_cell(b, y, x), layout.owner[y, x])

    def test_halo_pairs_are_owner_core_to_unique_halo(self):
        for ny, nx, by, bx, h in CASES:
            for periodic in PERIODIC:
                layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic)
                core = set(layout.owner.ravel().tolist())
                dst = layout.halo_dst.tolist()
                self.assertEqual(dst, sorted(dst))
                self.assertEqual(len(dst), len(set(dst)))
                self.assertTrue(core.isdisjoint(dst))
                self.assertTrue(set(layout.halo_src.tolist()) <= core)
                self.assertTrue(set(layout.halo_kind.tolist()) <= {FACE, CORNER})
        full = BlockLayout2D.build(6, 8, 3, 4)
        per_block = np.bincount(full.halo_kind, minlength=2) // full.nblocks
        np.testing.assert_array_equal(per_block, [2 * (3 + 4), 4])
        closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False))
        # Each 2x2 block keeps one inner face column, one face row and one corner.
        self.assertEqual(len(closed.halo_dst), 4 * (3 + 4 + 1))

    def test_pack_round_trip_and_exterior_stays_invalid(self):
        for ny, nx, by, bx, h in CASES:
            for periodic in PERIODIC:
                layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic)
                a = tags(ny, nx)
                blocked = layout.to_blocked(a)
                np.testing.assert_array_equal(layout.to_dense(blocked), a)
                filled = np.isfinite(blocked[:, 0]).sum()
                self.assertEqual(filled, ny * nx + len(layout.halo_dst))
                flat = blocked.ravel()
                for q in (0, 5):
                    np.testing.assert_array_equal(flat[layout.slot(layout.halo_dst, q).astype(np.intp)],
                                                  flat[layout.slot(layout.halo_src, q).astype(np.intp)])

    def test_tagged_pull_matches_dense_pull(self):
        for ny, nx, by, bx, h in CASES:
            for periodic in PERIODIC:
                with self.subTest(case=(ny, nx, by, bx, h), periodic=periodic):
                    layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic)
                    a = tags(ny, nx)
                    got = layout.to_dense(blocked_pull(layout, layout.to_blocked(a)))
                    np.testing.assert_array_equal(got, dense_pull(a, periodic))

    def test_block_cell_resolves_periodic_halo_images(self):
        layout = BlockLayout2D.build(6, 8, 3, 4)
        ly, lx = layout.local_shape
        # Block 0 sees x=7 (wrapped) in its west halo and y=5 in its south halo.
        self.assertEqual(layout.block_cell(0, 0, 7), 1 * lx + 0)
        self.assertEqual(layout.block_cell(0, 5, 0), 0 * lx + 1)
        closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False))
        with self.assertRaises(ValueError):
            closed.block_cell(0, 0, 7)
        with self.assertRaises(ValueError):
            layout.block_cell(0, 4, 0)  # two rows past the one-cell halo

    def test_invalid_inputs_rejected(self):
        for args, kw in (((6, 8, 4, 4), {}), ((6, 8, 3, 4), {"halo": 0}),
                         ((6, 8, 3, 4), {"halo": 4}), ((6, 8, 3, 4), {"periodic": (1, 1)}),
                         ((6, 8, 3, True), {}), ((6, 8.0, 3, 4), {})):
            with self.assertRaises(ValueError):
                BlockLayout2D.build(*args, **kw)
        layout = BlockLayout2D.build(6, 8, 3, 4)
        with self.assertRaises(ValueError):
            layout.to_blocked(np.zeros((NQ, 6, 8), dtype=np.int64))
        with self.assertRaises(ValueError):
            layout.to_dense(np.zeros((4, NQ, 5, 5)))

    def test_tables_are_readonly(self):
        layout = BlockLayout2D.build(6, 8, 3, 4)
        for array in (layout.origins, layout.owner, layout.halo_dst, layout.halo_src, layout.halo_kind):
            self.assertFalse(array.flags.writeable)


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