"""Standalone CPU manufactured checks for the oblique shear wave.

Run from this directory with PYTHONPATH=<checkout>/src (same as test_oracle.py).
"""

import hashlib
import io
import json
import math
import sys
import tempfile
from contextlib import redirect_stdout
from unittest.mock import patch
import unittest
from pathlib import Path

import numpy as np
import oracle
import run_case
import shear_wave as sw
from oracle import CX, CY, U, TIMES, centers, rectangles

N = 32
NU = U * N / sw.RE
HERE = Path(__file__).resolve().parent
ORACLE_SHA256_EEC4677 = "83e4d7425d45538278e9fb181af450a36358bd9a5ed6be57ff611dd9ae5eb35b"


def composite_masks(n):
    """Level-wide union masks for rectangles(n, 2): root complement + fine box."""
    (x0, y0, x1, y1), = rectangles(n, 2)[0]
    root = np.ones((n, n), dtype=bool)
    root[y0:y1, x0:x1] = False
    fine = np.zeros((2 * n, 2 * n), dtype=bool)
    fine[2 * y0:2 * y1, 2 * x0:2 * x1] = True
    return [root, fine]


def populations(masks, field, t=0.0):
    """Equilibrium G of an arbitrary (ux, uy, rho) field on every level."""
    out = []
    for level, m in enumerate(masks):
        x, y = centers(m.shape, level)
        ux, uy, rho = field(x, y, t)
        out.append(sw.equilibrium(rho, ux, uy))
    return out


def exact_field(phase=0.0, rate_factor=1.0, mode=sw.MODE, longitudinal=0.0, ripple=0.0,
                offset=0.0):
    k = 2 * np.pi / N * np.array(mode, dtype=float)
    e = np.array([-mode[1], mode[0]], dtype=float) / math.hypot(*mode)
    rate = rate_factor * sw.decay_rate(N, NU)
    kh = k / np.linalg.norm(k)

    def field(x, y, t):
        theta = k[0] * x + k[1] * y + phase
        a = U * np.exp(-rate * t)
        ux = a * e[0] * np.sin(theta) + longitudinal * kh[0] * np.sin(theta)
        uy = a * e[1] * np.sin(theta) + longitudinal * kh[1] * np.sin(theta)
        rho = 1 + offset + ripple * np.cos(theta)
        return ux, uy, rho

    return field


class ConventionTests(unittest.TestCase):
    def test_mode_angle_and_polarization(self):
        k = sw.wavevector(N)
        self.assertAlmostEqual(sw.ANGLE_DEG, 26.565051177077994)
        self.assertAlmostEqual(math.degrees(math.atan2(k[1], k[0])), sw.ANGLE_DEG)
        self.assertAlmostEqual(float(sw.E @ k), 0.0)
        self.assertAlmostEqual(float(sw.E @ sw.E), 1.0)
        self.assertAlmostEqual(sw.eps(128), 5 * (2 * math.pi / 128) ** 2)

    def test_navier_stokes_and_vanishing_advection(self):
        # Independent central differences: du/dt = nu lap u, (u.grad)u = 0, div u = 0.
        x = np.array([2.7, 8.1, 13.3])
        y = np.array([3.2, 7.4, 21.6])
        t, h = 10.0, 1e-3
        f = lambda x, y, t: np.array(sw.exact(x, y, t, N, NU)[:2])
        z = f(x, y, t)
        dt = (f(x, y, t + h) - f(x, y, t - h)) / (2 * h)
        dx = (f(x + h, y, t) - f(x - h, y, t)) / (2 * h)
        dy = (f(x, y + h, t) - f(x, y - h, t)) / (2 * h)
        lap = (f(x + h, y, t) + f(x - h, y, t) + f(x, y + h, t) + f(x, y - h, t) - 4 * z) / h**2
        np.testing.assert_allclose(dt, NU * lap, atol=1e-9)
        # Analytic identities: div u = A cos(theta) (E.k) and (u.grad)u = A sin(theta) (E.k) u/A.
        k = sw.wavevector(N)
        self.assertLess(abs(float(sw.E @ k)), 1e-17)
        # Central differences do not cancel exactly: each derivative carries
        # -(k_j h)^2/6 relative truncation, so div_h = -A cos(theta) h^2/6 sum_j E_j k_j^3
        # + O(h^4) (nonzero for an oblique k). Bound it with that leading term plus roundoff.
        a = U * math.exp(-sw.decay_rate(N, NU) * t)
        trunc = a * h * h / 6 * float(np.abs(sw.E * k**3).sum()) * 1.01
        # Argument rounding: eps*max|x|/h*U*|k| ~ 2.2e-16*21.6/1e-3*0.03*0.44 ~ 6e-14.
        roundoff = 1e-13
        self.assertLess(np.max(np.abs(dx[0] + dy[1])), trunc + roundoff)
        # Discrete advection: same truncation scaled by |u| <= a.
        self.assertLess(np.max(np.abs(z[0] * dx + z[1] * dy)), a * trunc + roundoff)
        # Halving h reduces the divergence truncation about fourfold.
        hh = h / 2
        div2 = ((f(x + hh, y, t) - f(x - hh, y, t))[0] + (f(x, y + hh, t) - f(x, y - hh, t))[1]) / (2 * hh)
        np.testing.assert_allclose(div2 * 4, dx[0] + dy[1], rtol=1e-2, atol=6 * roundoff)
        self.assertTrue(np.all(sw.exact(x, y, t, N, NU)[2] == 0))


class InitialTests(unittest.TestCase):
    def test_moments_on_each_level(self):
        for level, m in enumerate(composite_masks(N)):
            g = sw.initial(m.shape, level, N, NU)
            x, y = centers(m.shape, level)
            ux, uy, _ = sw.exact(x, y, 0.0, N, NU)
            rho = g.sum(axis=0)
            np.testing.assert_allclose(rho, 1.0, atol=1e-14)
            np.testing.assert_allclose(np.tensordot(CX, g, 1) / rho, ux, atol=1e-15)
            np.testing.assert_allclose(np.tensordot(CY, g, 1) / rho, uy, atol=1e-15)

    def test_equilibrium_matches_frozen_oracle_formula(self):
        # Same equilibrium as oracle.initial when fed the TG field.
        shape = (8, 8)
        x, y = centers(shape, 1)
        ux, uy, p = oracle.exact(x, y, 0.0, 4.0, 0.01)
        np.testing.assert_array_equal(
            sw.equilibrium(1 + 3 * p, ux, uy), oracle.initial(shape, 1, 4.0, 0.01)
        )


class MetricTests(unittest.TestCase):
    def setUp(self):
        self.masks = composite_masks(N)

    def sample(self, field, t):
        return sw.metrics(populations(self.masks, field, t), self.masks, t, N, NU, None)

    def test_exact_field_on_composite_partition(self):
        for t in (0.0, 100.0, 512.0):
            r = self.sample(exact_field(), t)
            self.assertLess(r["velocity_l2"], 1e-12)
            self.assertLess(r["residual_l2"], 1e-12)
            self.assertLess(r["longitudinal_l2"], 1e-12)
            self.assertLess(r["phase_drift_abs"], 1e-12)
            self.assertAlmostEqual(r["amplitude_ratio"], 1.0, places=12)
            self.assertLess(abs(r["amplitude_error"]), 1e-12)
            self.assertLess(r["spurious_pressure_scaled"], 1e-10)
            self.assertEqual(len(r["spurious_pressure_scaled_by_level"]), 2)
        self.assertEqual(r["leaf_cells"], int(sum(m.sum() for m in self.masks)))

    def test_phase_control(self):
        r = self.sample(exact_field(phase=0.01), 0.0)
        self.assertAlmostEqual(r["phase_drift"], 0.01, delta=1e-9)
        self.assertLess(abs(r["amplitude_error"]), 1e-12)  # R = hypot(a, b) is phase-free
        self.assertLess(r["residual_l2"], 1e-12)  # a pure phase shift stays modal

    def test_rate_control_and_fit(self):
        for factor, expected in ((1.0, 0.0), (1.01, 0.01)):
            field = exact_field(rate_factor=factor)
            rows = []
            for tstar in TIMES:
                step = round(tstar * N / U)
                rows.append(dict(self.sample(field, step), step=step))
            fit = sw.fit_rate(rows, N, NU)
            self.assertAlmostEqual(fit["rate_error"], expected, delta=1e-9)
            np.testing.assert_allclose(fit["interval_rate_errors"], expected, atol=1e-9)
            self.assertLess(fit["log_fit_residual_max"], 1e-9)
            self.assertAlmostEqual(fit["exact_rate_tstar"], 5 * (2 * math.pi) ** 2 / sw.RE)
            self.assertAlmostEqual(fit["exact_rate_tstar"], 1.97392088, places=8)

    def test_fit_ignores_startup_sample(self):
        rows = [dict(tstar=0.0, step=0, a_hat=2 * U, b_hat=0.0)]
        rows += [dict(tstar=s, step=round(s * N / U), b_hat=0.0,
                      a_hat=U * math.exp(-sw.decay_rate(N, NU) * round(s * N / U)))
                 for s in TIMES[1:]]
        self.assertLess(sw.fit_rate(rows, N, NU)["rate_error"], 1e-12)
        rows[-1]["a_hat"] = float("nan")
        with self.assertRaises(ValueError):
            sw.fit_rate(rows, N, NU)

    def test_wrong_mode_control(self):
        # Swapping (2,1) for (1,2), i.e. 26.6 vs 63.4 deg, is not silently modal.
        # Composite quadrature need not give an exactly zero projection.
        r = self.sample(exact_field(mode=(1, 2)), 0.0)
        self.assertLess(abs(r["amplitude_ratio"]), 0.05)
        self.assertGreater(r["residual_l2"], 0.95)

    def test_longitudinal_control(self):
        r = self.sample(exact_field(longitudinal=1e-3), 0.0)
        self.assertAlmostEqual(r["longitudinal_l2"], 1e-3 / U, delta=1e-9)
        # The full-vector residual and velocity error see it too.
        self.assertAlmostEqual(r["residual_l2"], 1e-3 / U, delta=1e-9)
        self.assertAlmostEqual(r["velocity_l2"], 1e-3 / U, delta=1e-9)
        self.assertLess(r["phase_drift_abs"], 1e-12)
        self.assertAlmostEqual(r["amplitude_ratio"], 1.0, places=12)

    def test_pressure_gauge_and_ripple(self):
        r = self.sample(exact_field(offset=1e-4), 0.0)
        self.assertLess(r["spurious_pressure_scaled"], 1e-10)  # one global gauge
        ripple = 3e-5
        r = self.sample(exact_field(ripple=ripple), 0.0)
        # p = ripple/3 cos(theta); rms over a whole number of periods = ripple/(3 sqrt 2).
        # Composite midpoint quadrature is not exactly periodic: 3% tolerance.
        expected = ripple / (3 * math.sqrt(2)) / (U * U / 4)
        self.assertAlmostEqual(r["spurious_pressure_scaled"], expected, delta=0.03 * expected)
        self.assertAlmostEqual(r["spurious_pressure_rms"], expected * U * U / 4,
                               delta=0.03 * expected * U * U / 4)

    def test_volume_weighting_and_double_count(self):
        bad = [np.ones_like(m) for m in self.masks]
        with self.assertRaises(ValueError):
            sw.metrics(populations(bad, exact_field()), bad, 0.0, N, NU)
        # Unweighted (count-based) averaging would change the fit; weights matter.
        r = self.sample(exact_field(), 0.0)
        self.assertEqual(r["leaf_cells"], N * N - self.masks[1].sum() // 4 + self.masks[1].sum())


def record(n, depth, rows, explosion="constant", completed=True, source="s"):
    for row in rows:
        row["spurious_pressure_scaled_by_level"] = [row["spurious_pressure_scaled"]] * depth
    return dict(n=n, depth=depth, viscosity=U * n / sw.RE, rows=rows, completed=completed,
                explosion=explosion, source=source)


def rows_for(n, rate_error=0.0, residual=1e-4, pressure=1e-3, phase=0.0):
    rate = (1 + rate_error) * sw.decay_rate(n, U * n / sw.RE)
    out = []
    for tstar in TIMES:
        step = round(tstar * n / U)
        amp = U * math.exp(-rate * step)
        out.append(dict(tstar=step * U / n, step=step, a_hat=amp, b_hat=0.0,
                        amplitude_ratio=1.0, modal_amplitude=amp, amplitude_exact=amp,
                        phase_drift=phase, mass=float(n * n), rho_min=0.999, rho_max=1.001,
                        max_speed=amp, longitudinal_l2=0.0, spurious_pressure_max_scaled=pressure,
                        amplitude_error=0.0, residual_l2=residual, phase_drift_abs=phase,
                        spurious_pressure_scaled=pressure,
                        spurious_pressure_rms=pressure * U * U / 4,
                        velocity_l2=residual, mass_relative=0.0))
    return out


class AssessTests(unittest.TestCase):
    def test_budget_form_is_max_not_sum(self):
        base = record(128, 1, rows_for(128, rate_error=1e-3, pressure=0.015))
        fine = record(256, 1, rows_for(256, rate_error=2e-4))
        # rate budget = max(2*1e-3, 0.002) = 0.002; 0.0025 would pass a sum form.
        refined = record(128, 2, rows_for(128, rate_error=0.0025))
        report = sw.assess(base, fine, [refined])
        self.assertEqual(report["status"], "FAIL")
        self.assertAlmostEqual(report["refined"]["constant"]["budget"]["rate_error"], 0.002)
        # pressure budget = max(2*0.015, 0.02) = 0.03 (control-dominated).
        self.assertAlmostEqual(
            report["refined"]["constant"]["budget"]["spurious_pressure_scaled"], 0.03)
        self.assertIn("finest", report["refined"]["constant"]["control_ratios"])
        ok = record(128, 2, rows_for(128, rate_error=0.0015), explosion="projected_linear")
        self.assertEqual(sw.assess(base, fine, [ok])["status"], "PASS")

    def test_control_outcomes_are_distinct(self):
        fine = record(256, 1, rows_for(256))
        refined = record(128, 2, rows_for(128))
        failing = record(128, 1, rows_for(128, residual=0.5))  # complete, > eps(128)
        report = sw.assess(failing, fine, [refined])
        self.assertEqual(report["status"], "CONTROL_SCREEN_FAIL")
        own = report["refined"]["constant"]
        self.assertTrue(own["own_screens_pass"])  # refined own screens kept
        self.assertFalse(own["controlled_pass"])  # but no controlled comparison passes
        good = record(128, 1, rows_for(128))
        self.assertEqual(sw.assess(None, fine, [refined])["status"], "INVALID_CONTROL")
        incomplete = record(128, 1, rows_for(128), completed=False)
        self.assertEqual(sw.assess(incomplete, fine, [refined])["status"], "INVALID_CONTROL")
        mixed = record(256, 1, rows_for(256), source="other")
        self.assertEqual(sw.assess(good, mixed, [refined])["status"], "INVALID_CONTROL")

    def test_invalid_rows_are_rejected_before_screens(self):
        base = record(128, 1, rows_for(128))
        fine = record(256, 1, rows_for(256))
        for mutate in (
            lambda rows: rows[2].update(rho_min=float("nan")),
            lambda rows: rows[2].update(spurious_pressure_scaled_by_level=[0.0, float("nan")]),
            lambda rows: rows[2].update(step=rows[1]["step"], tstar=rows[1]["tstar"]),
            lambda rows: rows[2].update(residual_l2=-1e-6),
            lambda rows: rows[2].pop("longitudinal_l2"),
        ):
            refined = record(128, 2, rows_for(128))
            mutate(refined["rows"])
            self.assertEqual(sw.assess(base, fine, [refined])["status"], "INVALID_RUN")
        # Legitimate signs stay allowed: negative amplitude and phase errors.
        refined = record(128, 2, rows_for(128))
        for row in refined["rows"]:
            row.update(amplitude_error=-1e-3, phase_drift=-1e-4, phase_drift_abs=1e-4)
        report = sw.assess(base, fine, [refined])
        self.assertEqual(report["status"], "PASS")
        self.assertEqual(sw.assess(base, fine, [])["status"], "INVALID_CONTROL")
        with self.assertRaises(ValueError):
            sw.assess(base, fine, [record(128, 2, rows_for(128)), record(128, 2, rows_for(128))])

    def test_modal_extinction_is_a_failure_not_malformed(self):
        def extinct(rows):
            for row in rows[1:]:
                row.update(a_hat=0.0, b_hat=0.0, modal_amplitude=0.0, amplitude_ratio=0.0,
                           amplitude_error=-1.0)
            return rows
        base = record(128, 1, rows_for(128))
        fine = record(256, 1, rows_for(256))
        refined = record(128, 2, extinct(rows_for(128)))
        report = sw.assess(base, fine, [refined])
        self.assertEqual(report["status"], "FAIL")
        own = report["refined"]["constant"]
        self.assertTrue(own["valid"])
        self.assertFalse(own["rate_fit_valid"])
        self.assertIsNone(own["rate_error"])
        self.assertIsNone(own["control_ratios"]["root"]["rate_error"])
        self.assertFalse(own["own_screens_pass"] or own["controlled_pass"] or own["budget_pass"])
        self.assertEqual(own["residual_l2"], 1e-4)  # other finite metrics retained
        json.loads(json.dumps(report, allow_nan=False))
        control = sw.assess(record(128, 1, extinct(rows_for(128))), fine,
                            [record(128, 2, rows_for(128))])
        self.assertEqual(control["status"], "CONTROL_SCREEN_FAIL")
        self.assertIsNone(control["refined"]["constant"]["budget"]["rate_error"])
        self.assertFalse(control["refined"]["constant"]["controlled_pass"])
        json.loads(json.dumps(control, allow_nan=False))

    def test_wrong_shapes(self):
        base = record(128, 1, rows_for(128))
        fine = record(256, 1, rows_for(256))
        refined = record(128, 2, rows_for(128))
        with self.assertRaises(ValueError):
            sw.assess(base, record(128, 1, rows_for(128)), [refined])
        with self.assertRaises(ValueError):
            sw.assess(record(128, 1, rows_for(128)), fine, [record(256, 2, rows_for(256))])

    def test_incomplete_record_fails(self):
        base = record(128, 1, rows_for(128))
        fine = record(256, 1, rows_for(256))
        refined = record(128, 2, rows_for(128), completed=False)
        self.assertEqual(sw.assess(base, fine, [refined])["status"], "FAIL")


class FakeRun:
    def __init__(self, arrays):
        self.arrays = arrays

    def download(self):
        return self.arrays if len(self.arrays) > 1 else self.arrays[0]


class RunnerSampleTests(unittest.TestCase):
    def test_refined_shear_sample_with_list_diagnostics(self):
        masks = composite_masks(N)
        arrays = [sw.initial(m.shape, level, N, NU) for level, m in enumerate(masks)]
        row = run_case.sample(FakeRun(arrays), masks, 0, N, NU, None, sw.metrics, 2)
        self.assertIsInstance(row["spurious_pressure_scaled_by_level"], list)
        self.assertTrue(run_case.finite_row(row))
        bad = [a.copy() for a in arrays]
        bad[1][0, masks[1]] = np.nan
        with self.assertRaises(FloatingPointError):
            run_case.sample(FakeRun(bad), masks, 0, N, NU, None, sw.metrics, 2, "READY")
        self.assertFalse(run_case.finite_row(dict(a=1.0, b=[0.0, float("inf")])))

    def test_taylor_green_sample_unchanged(self):
        masks = composite_masks(N)
        arrays = [oracle.initial(m.shape, level, N, NU) for level, m in enumerate(masks)]
        row = run_case.sample(FakeRun(arrays), masks, 0, N, NU, None, oracle.metrics, 2)
        self.assertEqual(row, oracle.metrics(arrays, masks, 0, N, NU, None))
        uniform = [oracle.initial((N, N), 0, N, NU)]
        row = run_case.sample(FakeRun(uniform), [np.ones((N, N), bool)], 0, N, NU, None,
                              oracle.metrics, 1)
        self.assertEqual(row["leaf_cells"], N * N)


def receipt(n, depth, explosion="constant", diagnostic=None):
    """Synthetic receipt satisfying every contract field of the fixed four-case set."""
    import assess_projected_pairs as tg
    steps = run_case.sample_steps(n)[0]
    rows = rows_for(n)
    for row, step in zip(rows, steps):
        row.update(step=step, tstar=step * U / n,
                   spurious_pressure_scaled_by_level=[row["spurious_pressure_scaled"]] * depth)
    cells = [n * n] if depth == 1 else [17 * n * n // 32, 15 * n * n // 8]
    hops = [16 * n * 2**level for level in range(depth)]
    builds = []
    for level in range(depth):
        size = n * 2**level
        builds.append(dict(source_sha256=tg.BUILD_SOURCES[size], arch="sm_70",
                           flags=["--use_fast_math"], n_objects=0, ptx_sha256="a" * 64,
                           config=dict(nx=10, ny=10, nz=1, lattice_model="D2Q9", collision="SRT",
                                       store_precision="FP32", streaming="double_buffer",
                                       ddf_shifting=True, les_enabled=False, velocity=0.0,
                                       viscosity=U * size / 100, rho=1.0,
                                       omega=1 / (.5 + 3 * U * size / 100),
                                       threads_per_block=256)))
    checkout = "/x/CelerisLab"
    return dict(
        n=n, depth=depth, re_L=100, explosion=explosion, U=U, mach=U * math.sqrt(3),
        viscosity=U * n / 100, tau_by_level=[0.5 + 3 * U * n / 100 * 2**l for l in range(depth)],
        source=dict(revision="r", dirty=False, source_tree_sha256=run_case.SOURCE_SHA256,
                    harness_sha256={"x.py": "h"}, build_capture_sha256=run_case.BUILD_CAPTURE_SHA256,
                    shear_wave_sha256=run_case.SHEAR_WAVE_SHA256),
        runtime_identity=dict(device=dict(name="V100", pci_bus_id="0", capability=[7, 0],
                                          driver=1, runtime=[12, 0]),
                              nvcc_path="/nvcc", nvcc_version="v", python="3", numpy="1",
                              pycuda="2"),
        runtime_source=dict(checkout=checkout, module=checkout + "/src/CelerisLab/__init__.py",
                            compiler=checkout + "/src/CelerisLab/cuda/compiler_v2.py",
                            kernel_root=checkout + "/src/CelerisLab/lbm/kernels"),
        cuda_visible_devices="0", numerical_revision=run_case.NUMERICAL_REVISION,
        reference_harness=run_case.REFERENCE_HARNESS, numerics=dict(tg.NUMERICS),
        initialization=sw.INITIALIZATION, limits=dict(sw.LIMITS),
        early_pressure=dict(enabled=False, rows=[], sampled_max_pressure_l2=None),
        refinements=[] if depth == 1 else rectangles(n, depth),
        workload=dict(active_cells_by_level=cells, level_hops=hops, root_intervals=16 * n,
                      active_collision_updates_by_level=[c * h for c, h in zip(cells, hops)],
                      active_collision_updates=sum(c * h for c, h in zip(cells, hops))),
        actual_build_sidecars=builds, rows=rows, completed=True,
        flow="shear_wave", mode=list(sw.MODE), polarization=sw.E.tolist(),
        angle_deg=sw.ANGLE_DEG, phase=sw.PHI, eps=sw.eps(n),
        **({"diagnostic": diagnostic} if diagnostic else {}))


class ContractTests(unittest.TestCase):
    def setUp(self):
        import assess_shear_wave
        self.assess = assess_shear_wave.assess
        self.records = [receipt(128, 1), receipt(256, 1), receipt(128, 2),
                        receipt(128, 2, "projected_linear")]

    def test_complete_contract_passes(self):
        result = self.assess(self.records, "r")
        self.assertTrue(result["contract_valid"], [k for k, v in result["checks"].items() if not v])
        self.assertEqual(result["status"], "PASS")

    def test_eps_metadata_is_portable_but_bounded(self):
        import assess_shear_wave
        self.assertTrue(assess_shear_wave.close_ulps(0.012047856934923534, 0.012047856934923532))
        exact = sw.eps(128)
        for value, ok in ((math.nextafter(exact, 1.0), True),
                          (math.nextafter(math.nextafter(exact, 0.0), 0.0), True),
                          (exact * (1 + 1e-12), False), (exact * 2, False),
                          (str(exact), False), (float("nan"), False), (None, False)):
            records = [dict(r) for r in self.records]
            records[2]["eps"] = value
            status = self.assess(records, "r")["status"]
            self.assertEqual(status == "PASS", ok, (value, status))

    def test_reproduced_false_passes_are_rejected(self):
        a = self.assess
        self.assertEqual(a(self.records[:2], "r")["status"], "INVALID_SET")  # empty refined
        dup = [*self.records[:3], receipt(128, 2)]
        self.assertEqual(a(dup, "r")["status"], "INVALID_SET")  # duplicate constant-only
        for mutate in (
            lambda r: r.pop("source"),
            lambda r: r.update(flow="taylor_green"),
            lambda r: r.update(mode=[1, 2]),
            lambda r: r.update(re_L=1000),
            lambda r: r.update(refinements=[[(0, 0, 8, 8)]]),
            lambda r: r["actual_build_sidecars"][0].update(arch="sm_80"),
            lambda r: r["rows"][0].update(step=1, tstar=U / 128),  # wrong initial time
            lambda r: r["source"].update(shear_wave_sha256="0" * 64),
        ):
            records = [dict(r, rows=[dict(x) for x in r["rows"]],
                            source=dict(r["source"]),
                            actual_build_sidecars=[dict(b) for b in r["actual_build_sidecars"]])
                       for r in self.records]
            mutate(records[2])
            self.assertEqual(a(records, "r")["status"], "INVALID_SET")
        mixed = [dict(r) for r in self.records]
        mixed[1] = dict(mixed[1], source=dict(mixed[1]["source"], revision="other"))
        self.assertEqual(a(mixed, "r")["status"], "INVALID_SET")
        self.assertEqual(a(self.records, "wrong-revision")["status"], "INVALID_SET")


def scaled_set(rate_error=0.0):
    d = run_case.SCALED_SHEAR
    out = [receipt(256, 1, diagnostic=d), receipt(512, 1, diagnostic=d),
           receipt(256, 2, diagnostic=d), receipt(256, 2, "projected_linear", diagnostic=d)]
    if rate_error:
        out[3]["rows"] = rows_for(256, rate_error=rate_error)
        for row, step in zip(out[3]["rows"], run_case.sample_steps(256)[0]):
            row.update(step=step, tstar=step * U / 256,
                       spurious_pressure_scaled_by_level=[row["spurious_pressure_scaled"]] * 2)
    return out


class ScaledContractTests(unittest.TestCase):
    def setUp(self):
        import assess_shear_wave
        self.mod = assess_shear_wave

    def test_scaled_contract_and_signed_excess(self):
        result = self.mod.assess(scaled_set(1.5e-3), "r", scaled=True)
        self.assertTrue(result["contract_valid"], [k for k, v in result["checks"].items() if not v])
        self.assertEqual([c[1] for c in result["cases"]], [256, 512, 256, 256])
        projected = result["signed"]["projected_linear"]
        self.assertAlmostEqual(projected["signed_rate_bias"], 1.5e-3, delta=1e-9)
        self.assertAlmostEqual(projected["excess_over_uniform_root"]["signed_rate_bias"], 1.5e-3,
                               delta=1e-9)
        self.assertIn("excess_over_uniform_fine", projected)

    def test_scaled_and_default_sets_do_not_mix(self):
        n128 = [receipt(128, 1), receipt(256, 1), receipt(128, 2), receipt(128, 2, "projected_linear")]
        self.assertEqual(self.mod.assess(n128, "r", scaled=True)["status"], "INVALID_SET")
        self.assertEqual(self.mod.assess(scaled_set(), "r")["status"], "INVALID_SET")
        unlabeled = scaled_set()
        unlabeled[0].pop("diagnostic")
        self.assertEqual(self.mod.assess(unlabeled, "r", scaled=True)["status"], "INVALID_SET")
        tg_label = scaled_set()
        tg_label[2]["diagnostic"] = run_case.SCALED_DIAGNOSTIC
        self.assertEqual(self.mod.assess(tg_label, "r", scaled=True)["status"], "INVALID_SET")

    def test_resolution_trend_and_reference_pin(self):
        reference = [receipt(128, 1), receipt(256, 1), receipt(128, 2),
                     receipt(128, 2, "projected_linear")]
        new = self.mod.assess(scaled_set(1.5e-3), "r", scaled=True)
        trend = self.mod.resolution_trend(new, reference, "r")
        self.assertTrue(trend["valid"])
        change = trend["changes"]["projected_linear"]["signed_rate_bias"]
        self.assertAlmostEqual(change["signed_change"], 1.5e-3, delta=1e-9)
        self.assertIn("excess_over_uniform_fine.residual_l2", trend["changes"]["constant"])
        with tempfile.TemporaryDirectory() as tmp:
            paths = []
            for i, r in enumerate(reference):
                paths.append(Path(tmp) / f"{i}.json")
                paths[-1].write_text(json.dumps(r))
            with self.assertRaises(ValueError):
                self.mod.load_reference(paths)

    def test_scaled_runner_options(self):
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            base = ["run_case.py", "--repo", str(repo), "--expected-revision", "f" * 40,
                    "--output", str(Path(tmp) / "case.json"), "--dry-run", "--re", "100"]
            shear = ["--flow", "shear_wave"]
            scaled = ["--diagnostic", run_case.SCALED_SHEAR]
            for extra, valid in (
                    (["--n", "512", "--depth", "1"] + shear + scaled, True),
                    (["--n", "256", "--depth", "2", "--explosion", "projected_linear"] + shear + scaled, True),
                    (["--n", "512", "--depth", "1"] + shear, False),
                    (["--n", "128", "--depth", "2"] + shear + scaled, False),
                    (["--n", "512", "--depth", "1"] + scaled, False),  # needs shear flow
                    (["--n", "512", "--depth", "1", "--early-pressure", "--diagnostic",
                      run_case.SCALED_DIAGNOSTIC] + shear, False),  # mutually exclusive
                    (["--n", "256", "--depth", "2", "--early-pressure"] + shear + scaled, False)):
                argv = base + extra
                with self.subTest(argv=argv), patch.object(sys, "argv", argv), \
                     patch.object(run_case, "verify_source", return_value={}), \
                     patch.object(run_case, "runtime_source", return_value={}), \
                     redirect_stdout(io.StringIO()), patch("sys.stderr", io.StringIO()):
                    if valid:
                        self.assertEqual(run_case.main(), 0)
                    else:
                        with self.assertRaises(SystemExit):
                            run_case.main()


class DefaultPathTests(unittest.TestCase):
    def test_frozen_oracle_is_byte_identical(self):
        digest = hashlib.sha256((HERE / "oracle.py").read_bytes()).hexdigest()
        self.assertEqual(digest, ORACLE_SHA256_EEC4677)

    def test_default_flow_is_frozen_taylor_green(self):
        init, met, limits, text = run_case.flow_functions()
        self.assertIs(init, oracle.initial)
        self.assertIs(met, oracle.metrics)
        self.assertIs(limits, oracle.LIMITS)
        self.assertEqual(text, run_case.INITIALIZATION)
        self.assertEqual(run_case.flow_functions("taylor_green"), (init, met, limits, text))
        shear = run_case.flow_functions("shear_wave")
        self.assertIs(shear[0], sw.initial)
        self.assertIs(shear[1], sw.metrics)
        with self.assertRaises(ValueError):
            run_case.flow_functions("tg")


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


class NestedPresetTests(unittest.TestCase):
    """Fixed three-level/two-region shear preset (nested_shear_n128)."""

    def setUp(self):
        import assess_shear_wave
        import shear_nested
        self.mod, self.nested = assess_shear_wave, shear_nested

    def test_exact_geometry_work_and_levels(self):
        nested = self.nested
        masks = nested.leaf_masks()
        self.assertEqual([int(m.sum()) for m in masks], [7936, 12288, 86016])
        self.assertEqual(nested.LEAVES, [7936, 12288, 86016])
        self.assertEqual(nested.HOPS, [2048, 4096, 8192])
        self.assertEqual(nested.WORK, 771227648)
        uniform = 128 * 128 * 2048 + 512 * 512 * 8192
        self.assertEqual(uniform + 2 * nested.WORK, 3723493376)
        np.testing.assert_allclose(nested.NU, [.0384, .0768, .1536], rtol=0, atol=1e-15)
        np.testing.assert_allclose(nested.TAU, [.6152, .7304, .9608], rtol=0, atol=1e-14)
        self.assertEqual(sum(float(m.sum()) * 0.25**l for l, m in enumerate(masks)), 128.0**2)
        # Production host hierarchy builds the same leaves from the same rectangles.
        from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition
        from CelerisLab.lbm.refinement import ACTIVE
        h = build_hierarchy(run_case.config(128, 100), nested.REFINEMENTS)
        np.testing.assert_array_equal(leaf_partition(h), np.ones((128, 128)))
        for mine, kinds in zip(masks, h.kinds):
            np.testing.assert_array_equal(mine, kinds == ACTIVE)

    def test_historical_shear_identity_is_preserved(self):
        digest = hashlib.sha256((HERE / "shear_wave.py").read_bytes()).hexdigest()
        self.assertEqual(digest, "81d5c8ee60864372887cddf9239580c82fb6750d956038f1e3ee9de61aa4e877")
        self.assertEqual(run_case.SHEAR_WAVE_SHA256, digest)
        nested_digest = hashlib.sha256((HERE / "shear_nested.py").read_bytes()).hexdigest()
        self.assertEqual(run_case.SHEAR_NESTED_SHA256, nested_digest)

    def nested_set(self):
        out = []
        for n, depth, scheme in ((128, 1, "constant"), (512, 1, "constant"),
                                 (128, 3, "constant"), (128, 3, "projected_linear")):
            r = receipt(n, depth, scheme, diagnostic=run_case.NESTED_SHEAR)
            r["source"]["shear_nested_sha256"] = run_case.SHEAR_NESTED_SHA256
            if depth == 3:
                cells, hops = self.nested.LEAVES, self.nested.HOPS
                r["refinements"] = [list(level) for level in self.nested.REFINEMENTS]
                r["workload"] = dict(active_cells_by_level=cells, level_hops=hops,
                                     root_intervals=2048,
                                     active_collision_updates_by_level=[c * h for c, h in zip(cells, hops)],
                                     active_collision_updates=self.nested.WORK)
            out.append(r)
        return out

    def test_nested_contract_and_rejections(self):
        a = self.mod.assess_nested
        good = self.nested_set()
        result = a(good, "r")
        self.assertTrue(result["contract_valid"], [k for k, v in result["checks"].items() if not v])
        self.assertEqual(result["status"], "PASS")
        self.assertIn("finest", result["physics"]["refined"]["constant"]["control_ratios"])
        wrong_finest = self.nested_set()
        wrong_finest[1] = receipt(256, 1, diagnostic=run_case.NESTED_SHEAR)
        wrong_finest[1]["source"]["shear_nested_sha256"] = run_case.SHEAR_NESTED_SHA256
        self.assertEqual(a(wrong_finest, "r")["status"], "INVALID_SET")
        for mutate in (
            lambda s: s[2].update(refinements=[s[2]["refinements"][1], s[2]["refinements"][0]]),
            lambda s: s[3].update(refinements=rectangles(128, 3)),
            lambda s: s[2]["rows"][2].update(spurious_pressure_scaled_by_level=[0.0, 0.0]),
            lambda s: s[0].update(diagnostic=run_case.SCALED_SHEAR),
            lambda s: s[2]["source"].pop("shear_nested_sha256"),
            lambda s: s[2].update(final_snapshot=dict(kept=True, sha256="a" * 64, bytes=10,
                                                      leaf_counts=self.nested.LEAVES, step=2048)),
            lambda s: s[3].update(final_snapshot=dict(kept=True, sha256="a" * 64,
                                                      bytes=self.nested.SNAPSHOT_MAX_BYTES + 1,
                                                      leaf_counts=self.nested.LEAVES, step=2048)),
        ):
            records = self.nested_set()
            mutate(records)
            self.assertNotEqual(a(records, "r")["status"], "PASS")
        ok = self.nested_set()
        ok[3]["final_snapshot"] = dict(kept=True, sha256="a" * 64, bytes=1000,
                                       leaf_counts=self.nested.LEAVES, step=2048)
        self.assertEqual(a(ok, "r")["status"], "PASS")
        with self.assertRaises(ValueError):
            self.nested.assess(good[0], good[1], [good[2], dict(good[2])])

    def test_nested_runner_options(self):
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            base = ["run_case.py", "--repo", str(repo), "--expected-revision", "f" * 40,
                    "--output", str(Path(tmp) / "case.json"), "--dry-run", "--re", "100"]
            nest = ["--flow", "shear_wave", "--diagnostic", run_case.NESTED_SHEAR]
            proj = ["--explosion", "projected_linear"]
            for extra, valid in (
                    (["--n", "128", "--depth", "1"] + nest, True),
                    (["--n", "512", "--depth", "1"] + nest, True),
                    (["--n", "128", "--depth", "3"] + nest, True),
                    (["--n", "128", "--depth", "3"] + proj + nest + ["--final-snapshot"], True),
                    (["--n", "256", "--depth", "1"] + nest, False),      # finest control is 4N
                    (["--n", "128", "--depth", "2"] + nest, False),
                    (["--n", "128", "--depth", "3", "--flow", "shear_wave"], False),
                    (["--n", "128", "--depth", "3", "--diagnostic", run_case.NESTED_SHEAR], False),
                    (["--n", "128", "--depth", "3"] + nest + ["--final-snapshot"], False),
                    (["--n", "128", "--depth", "3"] + proj + ["--flow", "shear_wave"], False),
                    (["--n", "128", "--depth", "2", "--flow", "shear_wave", "--final-snapshot"], False)):
                argv = base + extra
                with self.subTest(argv=argv), patch.object(sys, "argv", argv), \
                     patch.object(run_case, "verify_source", return_value={}) as verify, \
                     patch.object(run_case, "runtime_source", return_value={}), \
                     redirect_stdout(io.StringIO()), patch("sys.stderr", io.StringIO()):
                    if valid:
                        self.assertEqual(run_case.main(), 0)
                        self.assertEqual(verify.call_args.args[2], "shear_nested")
                    else:
                        with self.assertRaises(SystemExit):
                            run_case.main()

    def test_mocked_nested_execution_with_bounded_snapshot(self):
        from types import SimpleNamespace
        from contextlib import nullcontext
        from CelerisLab.lbm.hierarchy import build_hierarchy
        h = build_hierarchy(run_case.config(128, 100), self.nested.REFINEMENTS)
        nu = U * 128 / sw.RE
        arrays = [sw.initial(k.shape, level, 128, nu) for level, k in enumerate(h.kinds)]

        class FakeRun:
            def __init__(self, *args, **kwargs):
                self.step_count = 0
                self.kwargs = kwargs

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

            def initialize(self, *given):
                assert len(given) == 3

            def step(self, count):
                assert count >= 0
                self.step_count += count

            def download(self):
                return arrays

        modules = {"CelerisLab.cuda.context": SimpleNamespace(CudaContext=FakeRun),
                   "CelerisLab.lbm.block_stepper": SimpleNamespace(PeriodicBlockStepper=FakeRun),
                   "CelerisLab.lbm.refinement_stepper": SimpleNamespace(FixedPatchStepper=FakeRun)}
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            output = Path(tmp) / "projected_nested.json"
            argv = ["run_case.py", "--repo", str(repo), "--re", "100", "--n", "128", "--depth", "3",
                    "--explosion", "projected_linear", "--flow", "shear_wave",
                    "--diagnostic", run_case.NESTED_SHEAR, "--final-snapshot",
                    "--output", str(output), "--expected-revision", "f" * 40, "--run-gpu"]
            with patch.dict(sys.modules, modules), patch.object(sys, "argv", argv), \
                 patch.dict("os.environ", CUDA_VISIBLE_DEVICES="mock"), \
                 patch.object(run_case, "verify_source", return_value=dict(revision="f" * 40)), \
                 patch.object(run_case, "runtime_source", return_value={}), \
                 patch.object(run_case, "runtime_identity", return_value={}), \
                 patch.object(run_case, "build_capture", return_value=nullcontext()), \
                 redirect_stdout(io.StringIO()):
                run_case.main()
            rec = json.loads(output.read_text())
            self.assertTrue(rec["completed"])
            self.assertEqual(rec["diagnostic"], run_case.NESTED_SHEAR)
            self.assertEqual(rec["workload"]["active_cells_by_level"], self.nested.LEAVES)
            self.assertEqual(rec["workload"]["active_collision_updates"], self.nested.WORK)
            self.assertEqual(rec["tau_by_level"], self.nested.TAU)
            self.assertTrue(all(len(r["spurious_pressure_scaled_by_level"]) == 3 for r in rec["rows"]))
            snap = rec["final_snapshot"]
            path = output.with_name(output.stem + ".final_macro.npz")
            self.assertTrue(snap["kept"])
            self.assertEqual(snap["sha256"], hashlib.sha256(path.read_bytes()).hexdigest())
            self.assertLessEqual(snap["bytes"], self.nested.SNAPSHOT_MAX_BYTES)
            meta, fields = self.nested.read_snapshot(path)
            self.assertEqual((meta["step"], meta["phase"]), (2048, sw.PHI))
            self.assertEqual(meta["tstar"], 2048 * U / 128)
            self.assertEqual(meta["leaf_counts"], self.nested.LEAVES)
            self.assertEqual(sorted(fields), ["i", "j", "level", "rho", "ux", "uy"])
            np.testing.assert_allclose(fields["rho"], 1.0, atol=1e-12)

    def test_snapshot_contents_and_cap(self):
        nested = self.nested
        masks = nested.leaf_masks()
        arrays = [sw.initial(m.shape, level, 128, 0.0384) for level, m in enumerate(masks)]
        meta = nested.snapshot_meta(2048, 0.0384, "projected_linear", "r")
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "s.npz"
            block = nested.write_snapshot(path, arrays, masks, meta)
            self.assertTrue(block["kept"])
            got, fields = nested.read_snapshot(path)
            self.assertEqual(got["refinements"], [[list(r) for r in level] for level in nested.REFINEMENTS])
            self.assertEqual(int(fields["level"].max()), 2)
            with self.assertRaises(FileExistsError):
                nested.write_snapshot(path, arrays, masks, meta)
            bad = Path(tmp) / "bad.npz"
            np.savez_compressed(bad, populations=np.zeros(3), **{k: np.zeros(1) for k in ("level", "i", "j", "rho", "ux", "uy")}, meta=np.array("{}"))
            with self.assertRaises(ValueError):
                nested.read_snapshot(bad)
            with patch.object(nested, "SNAPSHOT_MAX_BYTES", 10):
                capped = nested.write_snapshot(Path(tmp) / "c.npz", arrays, masks, meta)
            self.assertFalse(capped["kept"])
            self.assertFalse((Path(tmp) / "c.npz").exists())


class TwinPresetTests(unittest.TestCase):
    """Depth-2 twin (twin_shear_n128): nested L0->L1 boxes, no L2, fresh U128/U256."""

    def setUp(self):
        import assess_shear_wave
        import shear_nested
        import shear_twin
        self.mod, self.nested, self.twin = assess_shear_wave, shear_nested, shear_twin

    def test_exact_geometry_work_and_nested_identity(self):
        twin = self.twin
        masks = twin.leaf_masks()
        self.assertEqual([int(m.sum()) for m in masks], [7936, 33792])
        self.assertEqual((twin.LEAVES, twin.HOPS, twin.WORK), ([7936, 33792], [2048, 4096], 154664960))
        self.assertEqual(128 * 128 * 2048 + 256 * 256 * 4096 + 2 * twin.WORK, 611319808)
        self.assertEqual(twin.REFINEMENTS, [self.nested.REFINEMENTS[0]])
        np.testing.assert_array_equal(masks[0], self.nested.leaf_masks()[0])  # same root leaves
        from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition
        from CelerisLab.lbm.refinement import ACTIVE
        h = build_hierarchy(run_case.config(128, 100), twin.REFINEMENTS)
        np.testing.assert_array_equal(leaf_partition(h), np.ones((128, 128)))
        for mine, kinds in zip(masks, h.kinds):
            np.testing.assert_array_equal(mine, kinds == ACTIVE)
        # Historical nested identity is unchanged; the twin is pinned separately.
        self.assertEqual(hashlib.sha256((HERE / "shear_nested.py").read_bytes()).hexdigest(),
                         "fa6cfe87fb4df8c6d732c637c885276833c81973b950c7724d734f8fc350bcb0")
        self.assertEqual(run_case.SHEAR_TWIN_SHA256,
                         hashlib.sha256((HERE / "shear_twin.py").read_bytes()).hexdigest())

    def twin_set(self):
        out = []
        for n, depth, scheme in ((128, 1, "constant"), (256, 1, "constant"),
                                 (128, 2, "constant"), (128, 2, "projected_linear")):
            r = receipt(n, depth, scheme, diagnostic=run_case.TWIN_SHEAR)
            r["source"].update(shear_nested_sha256=run_case.SHEAR_NESTED_SHA256,
                               shear_twin_sha256=run_case.SHEAR_TWIN_SHA256)
            if depth == 2:
                cells, hops = self.twin.LEAVES, self.twin.HOPS
                r["refinements"] = [list(level) for level in self.twin.REFINEMENTS]
                r["workload"] = dict(active_cells_by_level=cells, level_hops=hops, root_intervals=2048,
                                     active_collision_updates_by_level=[c * h for c, h in zip(cells, hops)],
                                     active_collision_updates=self.twin.WORK)
            out.append(r)
        return out

    def test_twin_contract_and_rejections(self):
        a = lambda records: self.mod.assess_nested(records, "r", run_case.TWIN_SHEAR)
        result = a(self.twin_set())
        self.assertTrue(result["contract_valid"], [k for k, v in result["checks"].items() if not v])
        self.assertEqual(result["status"], "PASS")
        self.assertIn("does not isolate abstract depth", result["scope"])
        self.assertEqual(result["finest_control"], "uniform 2N (U256)")
        for mutate in (
            lambda s: s.__setitem__(1, receipt(512, 1, diagnostic=run_case.TWIN_SHEAR)),  # 4N control
            lambda s: s[2].update(refinements=[list(level) for level in self.nested.REFINEMENTS]),
            lambda s: s[3].update(refinements=rectangles(128, 2)),  # old single box
            lambda s: s[2]["rows"][1].update(spurious_pressure_scaled_by_level=[0.0, 0.0, 0.0]),
            lambda s: s[0].update(diagnostic=run_case.NESTED_SHEAR),
            lambda s: s[2]["source"].pop("shear_twin_sha256"),
            lambda s: s[3].update(final_snapshot=dict(kept=True, sha256="a" * 64, bytes=10,
                                                      leaf_counts=self.twin.LEAVES, step=2048)),
        ):
            records = self.twin_set()
            mutate(records)
            self.assertNotEqual(a(records)["status"], "PASS")
        # Twin and nested sets never validate as each other.
        self.assertEqual(self.mod.assess_nested(self.twin_set(), "r")["status"], "INVALID_SET")

    def test_twin_runner_options(self):
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            base = ["run_case.py", "--repo", str(repo), "--expected-revision", "f" * 40,
                    "--output", str(Path(tmp) / "case.json"), "--dry-run", "--re", "100"]
            tw = ["--flow", "shear_wave", "--diagnostic", run_case.TWIN_SHEAR]
            proj = ["--explosion", "projected_linear"]
            for extra, valid in (
                    (["--n", "128", "--depth", "1"] + tw, True),
                    (["--n", "256", "--depth", "1"] + tw, True),
                    (["--n", "128", "--depth", "2"] + tw, True),
                    (["--n", "128", "--depth", "2"] + proj + tw, True),
                    (["--n", "512", "--depth", "1"] + tw, False),
                    (["--n", "128", "--depth", "3"] + tw, False),
                    (["--n", "128", "--depth", "2"] + proj + tw + ["--final-snapshot"], False),
                    (["--n", "128", "--depth", "2", "--diagnostic", run_case.TWIN_SHEAR], False)):
                argv = base + extra
                with self.subTest(argv=argv), patch.object(sys, "argv", argv), \
                     patch.object(run_case, "verify_source", return_value={}) as verify, \
                     patch.object(run_case, "runtime_source", return_value={}), \
                     redirect_stdout(io.StringIO()), patch("sys.stderr", io.StringIO()):
                    if valid:
                        self.assertEqual(run_case.main(), 0)
                        self.assertEqual(verify.call_args.args[2], "shear_twin")
                    else:
                        with self.assertRaises(SystemExit):
                            run_case.main()

    def test_mocked_twin_execution(self):
        from types import SimpleNamespace
        from contextlib import nullcontext
        from CelerisLab.lbm.hierarchy import build_hierarchy
        h = build_hierarchy(run_case.config(128, 100), self.twin.REFINEMENTS)
        arrays = [sw.initial(k.shape, level, 128, U * 128 / sw.RE) for level, k in enumerate(h.kinds)]
        selectors = []

        class FakeRun:
            def __init__(self, *args, **kwargs):
                self.step_count = 0
                if "refinements" in kwargs:
                    selectors.append((kwargs["explosion"], kwargs["refinements"]))

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

            def initialize(self, *given):
                assert len(given) == 2

            def step(self, count):
                self.step_count += count

            def download(self):
                return arrays

        modules = {"CelerisLab.cuda.context": SimpleNamespace(CudaContext=FakeRun),
                   "CelerisLab.lbm.block_stepper": SimpleNamespace(PeriodicBlockStepper=FakeRun),
                   "CelerisLab.lbm.refinement_stepper": SimpleNamespace(FixedPatchStepper=FakeRun)}
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            output = Path(tmp) / "projected_twin.json"
            argv = ["run_case.py", "--repo", str(repo), "--re", "100", "--n", "128", "--depth", "2",
                    "--explosion", "projected_linear", "--flow", "shear_wave",
                    "--diagnostic", run_case.TWIN_SHEAR, "--output", str(output),
                    "--expected-revision", "f" * 40, "--run-gpu"]
            with patch.dict(sys.modules, modules), patch.object(sys, "argv", argv), \
                 patch.dict("os.environ", CUDA_VISIBLE_DEVICES="mock"), \
                 patch.object(run_case, "verify_source", return_value=dict(revision="f" * 40)), \
                 patch.object(run_case, "runtime_source", return_value={}), \
                 patch.object(run_case, "runtime_identity", return_value={}), \
                 patch.object(run_case, "build_capture", return_value=nullcontext()), \
                 redirect_stdout(io.StringIO()):
                run_case.main()
            rec = json.loads(output.read_text())
            self.assertTrue(rec["completed"])
            self.assertEqual(rec["diagnostic"], run_case.TWIN_SHEAR)
            self.assertEqual(rec["workload"]["active_cells_by_level"], self.twin.LEAVES)
            self.assertEqual(rec["workload"]["active_collision_updates"], self.twin.WORK)
            self.assertNotIn("final_snapshot", rec)
            self.assertTrue(all(len(r["spurious_pressure_scaled_by_level"]) == 2 for r in rec["rows"]))
            self.assertEqual(selectors, [("projected_linear", [list(self.twin.REFINEMENTS[0])])])
            self.assertEqual(sorted(p.name for p in Path(tmp).iterdir()), ["projected_twin.json"])


class ProbeTraceTests(unittest.TestCase):
    """Signed interface-probe mechanism diagnostic (--probe-trace); no pass gates."""

    def setUp(self):
        import shear_nested
        import shear_probes
        import shear_twin
        self.nested, self.probes, self.twin = shear_nested, shear_probes, shear_twin

    def exact_arrays(self, masks, n, t):
        """Equilibrium of the exact wave at root-physical time t on a case grid of root N=n."""
        out = []
        for level, m in enumerate(masks):
            x, y = centers(m.shape, level)
            x, y = x * 128 / n, y * 128 / n
            ux, uy, _ = sw.exact(x, y, t, 128, U * 128 / sw.RE)
            out.append(sw.equilibrium(np.ones_like(ux), ux, uy))
        return out

    def test_frozen_history_identities(self):
        for name, digest in (("shear_wave.py", "81d5c8ee60864372887cddf9239580c82fb6750d956038f1e3ee9de61aa4e877"),
                             ("shear_nested.py", "fa6cfe87fb4df8c6d732c637c885276833c81973b950c7724d734f8fc350bcb0"),
                             ("shear_twin.py", "1cbf37f18e89aea577581631c4565f51e68d22a3739b698816ef64acbf2b2bb4"),
                             ("oracle.py", ORACLE_SHA256_EEC4677)):
            self.assertEqual(hashlib.sha256((HERE / name).read_bytes()).hexdigest(), digest, name)
        self.assertEqual(run_case.SHEAR_PROBES_SHA256,
                         hashlib.sha256((HERE / "shear_probes.py").read_bytes()).hexdigest())

    def test_probe_selection_per_case(self):
        sp = self.probes
        self.assertEqual(len(sp.points()), 24)
        corners = [(x, y) for level in (self.nested.REFINEMENTS[0],
                                        [tuple(v / 2 for v in r) for r in self.nested.REFINEMENTS[1]])
                   for x0, y0, x1, y1 in level for x in (x0, x1) for y in (y0, y1)]
        for _, _, (x, y) in sp.points():
            self.assertGreaterEqual(min(math.hypot(x - cx, y - cy) for cx, cy in corners), 8.0)
        pattern = [0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2]
        nested = sp.select(self.nested.leaf_masks())
        self.assertEqual([p["level"] for p in nested], pattern * 2)
        twin = sp.select(self.twin.leaf_masks())
        self.assertEqual([p["level"] for p in twin], [min(v, 1) for v in pattern] * 2)
        u128 = sp.select([np.ones((128, 128), bool)], 128)
        u512 = sp.select([np.ones((512, 512), bool)], 512)
        self.assertTrue(all(p["level"] == 0 for p in u128 + u512))
        for case in (nested, twin, u128, u512):
            self.assertEqual(len({(p["level"], p["i"], p["j"]) for p in case}), 24)  # unique
            for p in case:
                dx, dy = p["displacement"]
                self.assertEqual([p["center"][0] - p["point"][0], p["center"][1] - p["point"][1]],
                                 [dx, dy])
        self.assertEqual(max(math.hypot(*p["displacement"]) for p in u128), 0.625)
        self.assertLessEqual(max(math.hypot(*p["displacement"]) for p in u512), 0.25 * 0.7072)
        # Nested: 3 L0 + 5 L1 + 4 L2 per transect, as checked by native review.
        for name in sp.TRANSECTS:
            levels = [p["level"] for p in nested if p["transect"] == name]
            self.assertEqual([levels.count(v) for v in (0, 1, 2)], [3, 5, 4])
        # Different grids give different centers: no pretended collocation.
        self.assertNotEqual(u128[3]["center"], u512[3]["center"])
        self.assertEqual(sp.case_steps(512)[:3], [0, 4, 8])
        self.assertEqual(sp.case_steps(128), list(sp.STEPS))

    def test_residuals_vanish_for_the_exact_field(self):
        sp = self.probes
        for masks, n in ((self.nested.leaf_masks(), 128), ([np.ones((512, 512), bool)], 512)):
            probes = sp.select(masks, n)
            trace = sp.new_trace(probes, U * n / sw.RE, ["x", n, len(masks), "c"])
            for t in (0, 100, 2048):
                trace["steps"].append(t)
                trace["values"].append(sp.sample(self.exact_arrays(masks, n, t), probes))
            for res in sp.residuals(trace):
                self.assertLess(max(abs(v) for r in res for v in r), 1e-14)

    def test_validate_schedule_values_and_cap(self):
        sp = self.probes
        probes = sp.select(self.nested.leaf_masks())
        trace = sp.new_trace(probes, 0.0384, ["n", 128, 3, "projected_linear"])
        trace["steps"] = list(sp.STEPS)
        trace["values"] = [[1.0, 0.01, -0.02] * 24 for _ in sp.STEPS]
        size = sp.validate(trace)
        self.assertLess(size, sp.TRACE_MAX_BYTES)
        realistic = dict(trace, values=[[1.0000012345678901, 0.012345678901234567,
                                         -0.023456789012345678] * 24 for _ in sp.STEPS])
        self.assertLess(sp.validate(realistic), sp.TRACE_MAX_BYTES)
        for bad in (dict(trace, steps=trace["steps"][:-1]),
                    dict(trace, values=trace["values"][:-1] + [[float("nan")] * 72]),
                    dict(trace, probes=trace["probes"][:-1])):
            with self.assertRaises(ValueError):
                sp.validate(bad)
        with patch.object(sp, "TRACE_MAX_BYTES", 1000):
            with self.assertRaises(ValueError):
                sp.validate(trace)

    def test_probe_runner_options(self):
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            base = ["run_case.py", "--repo", str(repo), "--expected-revision", "f" * 40,
                    "--output", str(Path(tmp) / "case.json"), "--dry-run", "--re", "100",
                    "--flow", "shear_wave", "--probe-trace"]
            nest = ["--diagnostic", run_case.NESTED_SHEAR]
            tw = ["--diagnostic", run_case.TWIN_SHEAR]
            proj = ["--explosion", "projected_linear"]
            for extra, valid in (
                    (["--n", "128", "--depth", "1"] + nest, True),
                    (["--n", "512", "--depth", "1"] + nest, True),
                    (["--n", "128", "--depth", "3"] + proj + nest, True),
                    (["--n", "128", "--depth", "2"] + proj + tw, True),
                    (["--n", "128", "--depth", "3"] + nest, False),          # constant d3
                    (["--n", "128", "--depth", "1"] + tw, False),
                    (["--n", "128", "--depth", "3"] + proj + nest + ["--final-snapshot"], False),
                    (["--n", "128", "--depth", "2"], False)):
                argv = base + extra
                with self.subTest(argv=argv), patch.object(sys, "argv", argv), \
                     patch.object(run_case, "verify_source", return_value={}) as verify, \
                     patch.object(run_case, "runtime_source", return_value={}), \
                     redirect_stdout(io.StringIO()), patch("sys.stderr", io.StringIO()):
                    if valid:
                        self.assertEqual(run_case.main(), 0)
                        self.assertIs(verify.call_args.kwargs["probes"], True)
                    else:
                        with self.assertRaises(SystemExit):
                            run_case.main()

    def run_mocked(self, n, depth, diagnostic, scheme, masks):
        from types import SimpleNamespace
        from contextlib import nullcontext
        arrays = self.exact_arrays(masks, n, 0.0)
        counts = []

        class FakeRun:
            def __init__(self, *args, **kwargs):
                self.step_count = 0

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

            def initialize(self, *given):
                pass

            def step(self, count):
                self.step_count += count
                counts.append(self.step_count)

            def download(self):
                return arrays[0] if depth == 1 else arrays

        modules = {"CelerisLab.cuda.context": SimpleNamespace(CudaContext=FakeRun),
                   "CelerisLab.lbm.block_stepper": SimpleNamespace(PeriodicBlockStepper=FakeRun),
                   "CelerisLab.lbm.refinement_stepper": SimpleNamespace(FixedPatchStepper=FakeRun)}
        repo = Path(__file__).resolve().parents[3]
        with tempfile.TemporaryDirectory() as tmp:
            output = Path(tmp) / "probe.json"
            argv = ["run_case.py", "--repo", str(repo), "--re", "100", "--n", str(n),
                    "--depth", str(depth), "--explosion", scheme, "--flow", "shear_wave",
                    "--diagnostic", diagnostic, "--probe-trace", "--output", str(output),
                    "--expected-revision", "f" * 40, "--run-gpu"]
            with patch.dict(sys.modules, modules), patch.object(sys, "argv", argv), \
                 patch.dict("os.environ", CUDA_VISIBLE_DEVICES="mock"), \
                 patch.object(run_case, "verify_source", return_value=dict(revision="f" * 40)), \
                 patch.object(run_case, "runtime_source", return_value={}), \
                 patch.object(run_case, "runtime_identity", return_value={}), \
                 patch.object(run_case, "build_capture", return_value=nullcontext()), \
                 redirect_stdout(io.StringIO()):
                run_case.main()
            return json.loads(output.read_text()), counts, sorted(p.name for p in Path(tmp).iterdir())

    def test_mocked_probe_traces(self):
        sp = self.probes
        rec, counts, files = self.run_mocked(128, 3, run_case.NESTED_SHEAR, "projected_linear",
                                             self.nested.leaf_masks())
        trace = rec["probe_trace"]
        self.assertEqual(trace["steps"], list(sp.STEPS))
        self.assertEqual(len(trace["values"]), len(sp.STEPS))
        self.assertLessEqual(trace["bytes"], sp.TRACE_MAX_BYTES)
        self.assertEqual(trace["case"], [run_case.NESTED_SHEAR, 128, 3, "projected_linear"])
        self.assertEqual(len(rec["rows"]), 5)  # qualification samples preserved
        self.assertEqual(files, ["probe.json"])  # no field dumps
        self.assertEqual(sp.validate(trace) <= sp.TRACE_MAX_BYTES, True)
        rec512, counts512, _ = self.run_mocked(512, 1, run_case.NESTED_SHEAR, "constant",
                                               [np.ones((512, 512), bool)])
        t512 = rec512["probe_trace"]
        self.assertEqual(t512["steps"], list(sp.STEPS))  # root-physical times
        self.assertEqual(t512["viscosity_root"], U * 128 / sw.RE)
        self.assertEqual(t512["case_viscosity"], U * 512 / sw.RE)
        self.assertIn(4 * 1920, counts512)  # U512 advanced to 4x lattice steps
        self.assertTrue(all(p["level"] == 0 for p in t512["probes"]))
        res = sp.residuals(t512)
        self.assertLess(abs(res[0][0][1]), 1e-14)  # exact field at t=0 at actual centers