"""Independent unittest suite for the AMR 2:1 CPU volumetric reference (Grok, 2026-10-13).

Target: cpu_reference.py, API CPUREF-0.3 (also accepts CPU_REFERENCE_PATH env var).
Contract: AMR_Shared_CPU_API_2026-10-13.md, SHA256 ee8010e7...e112f6, plus
make_spec(*, collision='identity'|'srt').

Independence: expected values come from (a) a separate scalar-loop implementation
of the contract written here from the contract text and S-R Alg. 1 (it never uses
reference maps, masks, or helpers), and (b) a hand-derived lower-left SRT corner
oracle. Reference diagnostics are only checked against independently derived counts.
NumPy + stdlib only. <= 8 coarse steps per case.
"""
import importlib.util
import os
import sys
import unittest

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))


def _load_reference():
    cands = [os.environ.get("CPU_REFERENCE_PATH", ""),
             os.path.join(_HERE, "cpu_reference.py"),
             os.path.join(_HERE, "cpu_reference_v0.3.py"),
             os.path.join(_HERE, "cpu_reference_v0.2.py")]
    for p in cands:
        if p and os.path.exists(p):
            spec = importlib.util.spec_from_file_location("cpu_reference_under_test", p)
            mod = importlib.util.module_from_spec(spec)
            sys.modules[spec.name] = mod
            spec.loader.exec_module(mod)
            return mod, p
    raise FileNotFoundError("cpu_reference.py not found")


ref, REF_PATH = _load_reference()

# ---- independent constants from the contract -------------------------------
CX = (0, 1, -1, 0, 0, 1, -1, 1, -1)
CY = (0, 0, 0, 1, -1, 1, -1, -1, 1)
OPP = (0, 2, 1, 4, 3, 6, 5, 8, 7)
WT = (4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 36, 1 / 36)
N, NF, LO, HI, H = 16, 24, 4, 12, 4
TAU_C, TAU_F = 0.8, 1.1
# composite conservation tolerance: relative 1e-14 of |leaf total| over 8 steps (round-off only)
CONS_TOL = 1e-14


def covered(x, y):
    return LO <= x < HI and LO <= y < HI


def f_active(a, b):
    return H <= a < NF - H and H <= b < NF - H


def parent(a):
    return LO + (a - H) // 2


def f_inner(a, b):
    return (not f_active(a, b)) and LO - 1 <= parent(a) <= HI and LO - 1 <= parent(b) <= HI


def f_outer(a, b):
    return not f_active(a, b) and not f_inner(a, b)


def children(kx, ky):
    a0, b0 = H + 2 * (kx - LO), H + 2 * (ky - LO)
    return [(a, b) for b in (b0, b0 + 1) for a in (a0, a0 + 1)]  # (x, y)


def feq_point(rho, ux, uy):
    out = []
    for q in range(9):
        cu = CX[q] * ux + CY[q] * uy
        out.append(WT[q] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy)))
    return out


def srt_point(f, tau):
    rho = sum(f)
    ux = sum(CX[q] * f[q] for q in range(9)) / rho
    uy = sum(CY[q] * f[q] for q in range(9)) / rho
    e = feq_point(rho, ux, uy)
    return [f[q] - (f[q] - e[q]) / tau for q in range(9)]


def slot_kind(q, kx, ky):
    """Coalescence slot of non-covered coarse K, direction q, from link geometry only."""
    if q == 0 or covered(kx, ky):
        return None
    sx, sy = (kx - CX[q]) % N, (ky - CY[q]) % N
    if covered(sx, sy):
        return "face_or_diag_from_patch"
    if CX[q] != 0 and CY[q] != 0:
        cxp, cyp = kx + 0.5 - CX[q] / 2, ky + 0.5 - CY[q] / 2  # shared corner point
        if cxp in (LO, HI) and cyp in (LO, HI):
            return "corner_pass"
    return None


def independent_step(fc, ff, srt):
    """Scalar-loop implementation of one synchronised coarse step (contract + S-R Alg. 1)."""
    col = (lambda v, t: srt_point(v, t)) if srt else (lambda v, t: list(v))
    gc = np.full_like(fc, np.nan)
    for y in range(N):
        for x in range(N):
            if not covered(x, y):
                gc[:, y, x] = col(fc[:, y, x], TAU_C)
    gf = np.full_like(ff, np.nan)
    for b in range(NF):
        for a in range(NF):
            if f_active(a, b):
                gf[:, b, a] = col(ff[:, b, a], TAU_F)
            else:  # explosion into every ghost, once
                gf[:, b, a] = gc[:, parent(b), parent(a)]

    def hop(g):
        out = np.full_like(g, np.nan)
        for b in range(NF):
            for a in range(NF):
                if f_outer(a, b):
                    continue
                for q in range(9):
                    sa, sb = a - CX[q], b - CY[q]
                    if 0 <= sa < NF and 0 <= sb < NF:
                        out[q, b, a] = g[q, sb, sa]
        return out

    f1 = hop(gf)
    g1 = np.full_like(f1, np.nan)
    for b in range(NF):
        for a in range(NF):
            if f_active(a, b):
                g1[:, b, a] = col(f1[:, b, a], TAU_F)
            elif f_inner(a, b):
                g1[:, b, a] = f1[:, b, a]
    f2 = hop(g1)
    fc_new = np.full_like(fc, np.nan)
    for y in range(N):
        for x in range(N):
            if covered(x, y):
                continue
            for q in range(9):
                if slot_kind(q, x, y):
                    vals = [f2[q, b, a] for a, b in children(x, y)]
                    fc_new[q, y, x] = sum(vals) / 4.0
                else:
                    fc_new[q, y, x] = gc[q, (y - CY[q]) % N, (x - CX[q]) % N]
    ff_new = np.full_like(ff, np.nan)
    for b in range(NF):
        for a in range(NF):
            if f_active(a, b):
                ff_new[:, b, a] = f2[:, b, a]
    return fc_new, ff_new


def leaf(fc, ff):
    tot = np.zeros(3)
    for y in range(N):
        for x in range(N):
            if not covered(x, y):
                for q in range(9):
                    tot += fc[q, y, x] * np.array([1.0, CX[q], CY[q]])
    for b in range(NF):
        for a in range(NF):
            if f_active(a, b):
                for q in range(9):
                    tot += 0.25 * ff[q, b, a] * np.array([1.0, CX[q], CY[q]])
    return tot


def perturbed_state(spec, seed, amp=1e-3, u=(0.03, -0.02)):
    st = ref.make_state(spec, rho=1.0, velocity=u)
    rng = np.random.default_rng(seed)
    fc, ff = st.coarse_f.copy(), st.fine_f.copy()
    fc = np.where(np.isfinite(fc), fc * (1 + amp * rng.standard_normal(fc.shape)), np.nan)
    ff = np.where(np.isfinite(ff), ff * (1 + amp * rng.standard_normal(ff.shape)), np.nan)
    return ref.State(coarse_f=fc, fine_f=ff, phase="arrival_f", tick=st.tick)


def run(spec, st, steps):
    maps = ref.build_maps(spec)
    diags = []
    for _ in range(steps):
        st, d = ref.advance_coarse_once(spec, maps, st)
        diags.append(d)
    return st, diags


class TestIdentityTransportConservation(unittest.TestCase):
    def test_leaf_mass_momentum_8_steps(self):
        spec = ref.make_spec(collision="identity")
        st = perturbed_state(spec, 1)
        m0 = leaf(st.coarse_f, st.fine_f)
        st8, diags = run(spec, st, 8)
        m8 = leaf(st8.coarse_f, st8.fine_f)
        print(f"\n[identity] leaf0={m0.tolist()} leaf8={m8.tolist()} drift={(m8 - m0).tolist()}")
        np.testing.assert_allclose(m8, m0, rtol=0, atol=CONS_TOL * np.abs(m0).max())
        for d in diags:
            np.testing.assert_allclose(d["budget"]["interface_residual"], 0, atol=1e-12)
            np.testing.assert_allclose(d["leaf_after"] - d["leaf_before"], 0, atol=1e-12)
        self.assertEqual(st8.tick, 16)

    def test_matches_independent_implementation(self):
        spec = ref.make_spec(collision="identity")
        st = perturbed_state(spec, 2)
        fc, ff = st.coarse_f.copy(), st.fine_f.copy()
        maps = ref.build_maps(spec)
        worst = 0.0
        for _ in range(3):
            st, _ = ref.advance_coarse_once(spec, maps, st)
            fc, ff = independent_step(fc, ff, srt=False)
            np.testing.assert_array_equal(np.isnan(st.coarse_f), np.isnan(fc))
            np.testing.assert_array_equal(np.isnan(st.fine_f), np.isnan(ff))
            worst = max(worst, np.nanmax(np.abs(st.coarse_f - fc)), np.nanmax(np.abs(st.fine_f - ff)))
        print(f"\n[identity] max |ref - independent| over 3 steps = {worst:.3e}")
        self.assertLessEqual(worst, 1e-15)


class TestSRT(unittest.TestCase):
    def test_matches_independent_implementation(self):
        spec = ref.make_spec(collision="srt")
        st = perturbed_state(spec, 3)
        fc, ff = st.coarse_f.copy(), st.fine_f.copy()
        maps = ref.build_maps(spec)
        worst = 0.0
        for _ in range(3):
            st, _ = ref.advance_coarse_once(spec, maps, st)
            fc, ff = independent_step(fc, ff, srt=True)
            np.testing.assert_array_equal(np.isnan(st.coarse_f), np.isnan(fc))
            worst = max(worst, np.nanmax(np.abs(st.coarse_f - fc)), np.nanmax(np.abs(st.fine_f - ff)))
        print(f"\n[srt] max |ref - independent| over 3 steps = {worst:.3e}")
        self.assertLessEqual(worst, 1e-14)

    def test_leaf_conservation_and_collision_moments_8_steps(self):
        spec = ref.make_spec(collision="srt")
        st = perturbed_state(spec, 4)
        m0 = leaf(st.coarse_f, st.fine_f)
        st8, diags = run(spec, st, 8)
        m8 = leaf(st8.coarse_f, st8.fine_f)
        print(f"\n[srt] leaf0={m0.tolist()} leaf8={m8.tolist()} drift={(m8 - m0).tolist()}")
        np.testing.assert_allclose(m8, m0, rtol=0, atol=CONS_TOL * np.abs(m0).max())
        for d in diags:
            np.testing.assert_allclose(d["budget"]["coarse_collision"], 0, atol=1e-12)
            np.testing.assert_allclose(d["budget"]["fine_collision"], 0, atol=1e-12)
            np.testing.assert_allclose(d["budget"]["interface_residual"], 0, atol=1e-12)
        # independent per-cell check of SRT moment conservation
        rng = np.random.default_rng(9)
        f = np.array(feq_point(1.0, 0.04, -0.01)) * (1 + 1e-2 * rng.standard_normal(9))
        g = srt_point(list(f), TAU_F)
        for w in ([1.0] * 9, CX, CY):
            self.assertAlmostEqual(float(np.dot(w, g)), float(np.dot(w, f)), places=14)


class TestOwnershipPoison(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.spec = ref.make_spec(collision="srt")
        cls.maps = ref.build_maps(cls.spec)
        cls.st = perturbed_state(cls.spec, 5)
        cls.out, cls.d = ref.advance_coarse_once(cls.spec, cls.maps, cls.st)

    def test_collision_and_explosion_counts(self):
        d = self.d
        exp_c = np.array([[0 if covered(x, y) else 1 for x in range(N)] for y in range(N)])
        exp_f = np.array([[2 if f_active(a, b) else 0 for a in range(NF)] for b in range(NF)])
        exp_e = np.array([[0 if f_active(a, b) else 1 for a in range(NF)] for b in range(NF)])
        np.testing.assert_array_equal(d["collision_count_coarse"], exp_c)
        np.testing.assert_array_equal(d["collision_count_fine"], exp_f)
        np.testing.assert_array_equal(d["explosion_writes"], exp_e)
        self.assertNotIn("explosion_count", d)
        self.assertNotIn("fine_valid_after", d)
        self.assertIn("hop2_valid", d)
        self.assertNotIn("closure_error", d["budget"])
        self.assertIn("bookkeeping_check", d["budget"])
        # invalid tick rejection (v0.3)
        for bad in (1, -2, 1.0, True, 2.5, "0"):
            with self.assertRaises(ValueError):
                ref.advance_coarse_once(self.spec, self.maps,
                    ref.State(coarse_f=self.st.coarse_f, fine_f=self.st.fine_f, phase="arrival_f", tick=bad))
        out4, d4 = ref.advance_coarse_once(self.spec, self.maps,
            ref.State(coarse_f=self.st.coarse_f, fine_f=self.st.fine_f, phase="arrival_f", tick=4))
        self.assertEqual((d4["tick_before"], d4["tick_after"], out4.tick), (4, 6, 6))
        print(f"\ncounts: coarse cells collided once={int(exp_c.sum())}, fine twice={int((exp_f == 2).sum())}, "
              f"ghost explosion writes={int(exp_e.sum())}; tick validation OK")

    def test_coalescence_slots_reads_owner(self):
        d = self.d
        slots = {(q, y, x): slot_kind(q, x, y) for y in range(N) for x in range(N) for q in range(9)
                 if slot_kind(q, x, y)}
        n_pass = sum(1 for v in slots.values() if v == "corner_pass")
        got = {tuple(int(v) for v in r) for r in np.argwhere(d["coalesce_writes"] > 0)}
        self.assertEqual(got, set(slots))
        self.assertTrue((d["coalesce_writes"][d["coalesce_writes"] > 0] == 1).all())
        exp_reads = np.zeros((9, NF, NF), int)
        for (q, y, x) in slots:
            for a, b in children(x, y):
                self.assertTrue(f_inner(a, b))
                exp_reads[q, b, a] += 1
        np.testing.assert_array_equal(d["coalesce_reads"], exp_reads)
        own = d["coarse_incoming_owner"]
        n_cov = sum(covered(x, y) for x in range(N) for y in range(N))
        self.assertEqual(int((own == 2).sum()), len(slots))
        self.assertEqual(int((own == 0).sum()), 9 * n_cov)
        self.assertEqual(int((own == 1).sum()), 9 * (N * N - n_cov) - len(slots))
        print(f"\nslots={len(slots)} (corner-pass {n_pass}), unique reads={int((exp_reads == 1).sum())}, "
              f"max read multiplicity={int(exp_reads.max())}")

    def test_hop_sources_never_outer(self):
        for h in ("hop1", "hop2"):
            s = self.d["fine_hop_sources"][h]
            self.assertEqual(s["outer_ghost"], 0)
            self.assertEqual(s["outside"], 0)
            self.assertEqual(sum(s.values()), 9 * 256)

    def test_invalid_input_entries_ignored_and_poisoned(self):
        fc, ff = self.st.coarse_f.copy(), self.st.fine_f.copy()
        fc[:, LO:HI, LO:HI] = 1e300
        for b in range(NF):
            for a in range(NF):
                if not f_active(a, b):
                    ff[:, b, a] = -7.0
        st2 = ref.State(coarse_f=fc, fine_f=ff, phase="arrival_f", tick=0)
        out2, _ = ref.advance_coarse_once(self.spec, self.maps, st2)
        np.testing.assert_array_equal(out2.coarse_f, self.out.coarse_f)
        np.testing.assert_array_equal(out2.fine_f, self.out.fine_f)
        self.assertTrue(np.isnan(self.out.coarse_f[:, LO:HI, LO:HI]).all())
        ghost = np.array([[not f_active(a, b) for a in range(NF)] for b in range(NF)])
        self.assertTrue(np.isnan(self.out.fine_f[:, ghost]).all())
        self.assertTrue(np.isfinite(self.out.fine_f[:, ~ghost]).all())

    def test_inputs_not_mutated_and_maps_readonly(self):
        spec = self.spec
        maps = ref.build_maps(spec)
        st = perturbed_state(spec, 6)
        c0, f0 = st.coarse_f.copy(), st.fine_f.copy()
        ref.advance_coarse_once(spec, maps, st)
        np.testing.assert_array_equal(st.coarse_f, c0)
        np.testing.assert_array_equal(st.fine_f, f0)
        for name in ("coarse_active", "fine_active", "fine_inner_ghost", "coalesce"):
            with self.assertRaises(ValueError):
                getattr(maps, name)[...] = 0

    def test_nonfinite_active_input_rejected(self):
        fc = self.st.coarse_f.copy()
        fc[3, 0, 0] = np.nan
        with self.assertRaises(ValueError):
            ref.advance_coarse_once(self.spec, self.maps,
                                    ref.State(coarse_f=fc, fine_f=self.st.fine_f, phase="arrival_f", tick=0))

    def test_make_spec_api(self):
        self.assertEqual(ref.make_spec().collision, "identity")
        self.assertEqual(ref.make_spec(collision="srt").collision, "srt")
        with self.assertRaises(ValueError):
            ref.make_spec(collision="mrt")


# corner-pass links: (name, q, source (x,y), target K (x,y)); 8 total, two per corner
CORNER_PASS = [
    ("LL_8", 8, (4, 3), (3, 4)), ("LL_7", 7, (3, 4), (4, 3)),
    ("LR_5", 5, (11, 3), (12, 4)), ("LR_6", 6, (12, 4), (11, 3)),
    ("UL_6", 6, (4, 12), (3, 11)), ("UL_5", 5, (3, 11), (4, 12)),
    ("UR_7", 7, (11, 12), (12, 11)), ("UR_8", 8, (12, 11), (11, 12)),
]


class TestFourCorners(unittest.TestCase):
    def test_corner_pass_links_identity_marker(self):
        """Identity: a marker in the source's post-collision slot q must arrive exactly at K."""
        spec = ref.make_spec(collision="identity")
        maps = ref.build_maps(spec)
        for name, q, (sx, sy), (kx, ky) in CORNER_PASS:
            self.assertEqual(slot_kind(q, kx, ky), "corner_pass", name)
            st = ref.make_state(spec)
            fc = st.coarse_f.copy()
            fc[q, sy, sx] += 0.01
            out, d = ref.advance_coarse_once(spec, maps, ref.State(fc, st.fine_f, "arrival_f", 0))
            self.assertEqual(int(d["coarse_incoming_owner"][q, ky, kx]), 2, name)
            self.assertAlmostEqual(out.coarse_f[q, ky, kx], fc[q, sy, sx], places=15, msg=name)
        listed = {(link.key, link.i, link.source_xy, link.target_xy) for link in maps.corner_links}
        self.assertEqual(len(listed), 8)
        self.assertEqual(len({link.key for link in maps.corner_links}), 8)
        print(f"\nreference corner_links lists {len(listed)} uniquely keyed corner-pass links")

    def test_four_corner_symmetry_srt(self):
        """Mirror the lower-left SRT corner case to the other three corners; values must match."""
        spec = ref.make_spec(collision="srt")
        maps = ref.build_maps(spec)
        eps = 1e-3
        flips = {"LL": (False, False), "LR": (True, False), "UL": (False, True), "UR": (True, True)}
        mirror_q = {(True, False): (0, 2, 1, 3, 4, 8, 7, 6, 5), (False, True): (0, 1, 2, 4, 3, 7, 8, 5, 6)}

        def mq(q, fx, fy):
            if fx:
                q = mirror_q[(True, False)][q]
            if fy:
                q = mirror_q[(False, True)][q]
            return q

        vals = {}
        for tag, (fx, fy) in flips.items():
            st = ref.make_state(spec)
            fc = st.coarse_f.copy()
            sx, sy = (N - 1 - 4 if fx else 4), (N - 1 - 3 if fy else 3)
            for q, s in ((5, 1), (6, 1), (7, -1), (8, -1)):
                fc[mq(q, fx, fy), sy, sx] += s * eps
            out, _ = ref.advance_coarse_once(spec, maps, ref.State(fc, st.fine_f, "arrival_f", 0))
            kx, ky = (N - 1 - 3 if fx else 3), (N - 1 - 4 if fy else 4)
            vals[tag] = out.coarse_f[mq(8, fx, fy), ky, kx]
        print(f"\ncorner SRT symmetry values: {vals}")
        for tag in vals:
            self.assertAlmostEqual(vals[tag], vals["LL"], places=15, msg=tag)


class TestSRTCornerOracle(unittest.TestCase):
    """Contract oracle: LL corner, coarse source (4.5,3.5) -> target (3.5,4.5), q=8,
    source perturbed by eps*(e5+e6-e7-e8). Expected value computed by hand, no reference maps."""

    def _oracle(self, eps, srt):
        base = feq_point(1.0, 0.0, 0.0)
        src = list(base)
        for q, s in ((5, 1), (6, 1), (7, -1), (8, -1)):
            src[q] += s * eps
        g_src = srt_point(src, TAU_C) if srt else src  # coarse collision at (4,3)
        # Children of coarse (4,3): fine (4,2),(5,2),(4,3),(5,3). Only (5,3)->(4,4) hits an
        # active intermediate; the other three pass inner ghosts (3,3),(4,3),(3,4) uncollided.
        # Arrival at active (4,4) after hop 1: q=8 from (5,3) [parent (4,3)] carries g_src[8];
        # q=3 from (4,3) [parent (4,3)] carries g_src[3] (unperturbed); every other direction
        # comes from unperturbed equilibrium cells (active or ghost of unperturbed parents).
        arr = list(base)
        arr[8] = g_src[8]
        arr[3] = g_src[3]
        mid = srt_point(arr, TAU_F) if srt else arr
        return (3 * g_src[8] + mid[8]) / 4.0, g_src[8]

    def _run(self, collision, eps):
        spec = ref.make_spec(collision=collision)
        st = ref.make_state(spec)
        fc = st.coarse_f.copy()
        for q, s in ((5, 1), (6, 1), (7, -1), (8, -1)):
            fc[q, 3, 4] += s * eps
        out, d = ref.advance_coarse_once(spec, ref.build_maps(spec), ref.State(fc, st.fine_f, "arrival_f", 0))
        return out.coarse_f[8, 4, 3], d

    def test_srt_oracle(self):
        eps = 1e-3
        exp, g8 = self._oracle(eps, True)
        got, d = self._run("srt", eps)
        print(f"\n[srt corner] expected={exp!r} got={got!r} diff={got - exp:.3e} (uncoalesced g8={g8!r})")
        self.assertAlmostEqual(got, exp, places=15)
        self.assertNotAlmostEqual(exp, g8, places=9)  # mid-step collision really matters
        self.assertAlmostEqual(d["corner_packets"]["lower_left_q8"]["coalesced"], exp, places=15)

    def test_identity_changes_expected_packet(self):
        eps = 1e-3
        exp_srt, _ = self._oracle(eps, True)
        exp_id, _ = self._oracle(eps, False)
        got_id, _ = self._run("identity", eps)
        print(f"\n[identity corner] expected={exp_id!r} got={got_id!r}; srt-expected differs by {exp_srt - exp_id:.3e}")
        self.assertAlmostEqual(got_id, exp_id, places=15)
        self.assertAlmostEqual(exp_id, 1 / 36 - eps, places=15)
        self.assertGreater(abs(exp_srt - exp_id), 1e-6)


if __name__ == "__main__":
    import platform
    import resource
    import sys
    import hashlib
    print("reference:", REF_PATH, "API", ref.API_VERSION)
    for p in (REF_PATH, os.path.abspath(__file__)):
        print("sha256", hashlib.sha256(open(p, "rb").read()).hexdigest(), os.path.basename(p))
    print("python", sys.version.split()[0], "numpy", np.__version__, platform.platform())
    prog = unittest.main(exit=False, verbosity=2)
    print("peak_rss_MiB", resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024)
    sys.exit(0 if prog.result.wasSuccessful() else 1)