"""Optional observation-window exchange ledger: host contract (mocked driver) and GPU gate.

Run CPU and GPU classes in separate processes; GPU callers select the device.
ExchangeWindowGPU uses the SRT/unshifted Jet fixture; ExchangeWindowShiftedMRTGPU covers a
shifted MRT, all-blocked, stationary-body case without Jets (no shifted-MRT Jet claim).
GPU tolerance (conservative engineering tolerance, not a universal proof): the window ledger
sums h*dt*raw per link over fine substeps, while the reference sums per-base interval ledgers
over links and then over base steps. Both are FP64 reorderings of the same FP32-exact substep
contributions, so the scale is the sum of the ACTUAL absolute substep contributions (read from
last_raw after every fine hop), not of per-base totals that could hide substep cancellation:
|window - reference| <= 4 * m * eps * sum|substep terms| + 1e-300, m = links * substeps.
"""
from contextlib import ExitStack
import unittest
from unittest.mock import patch

import numpy as np
import pycuda.driver as cuda

from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from test_refinement_blocked_base import mocked_backend
from test_refinement_blocked_fine import BOUNDS, circle, config

EPS = np.finfo(np.float64).eps


def build(m, window=True, raw=True):
    return FixedPatchStepper(config("periodic"), context=m["context"], refinements=[BOUNDS],
                             level_blocks=[(3, 4), (16, 20)], level_objects=[(), (circle(),)],
                             observe_raw_exchange=raw, observe_window_exchange=window)


def same_bits(test, a, b):
    """Identical NaN masks, no infinities, identical finite FP32 bits (incl. signed zero)."""
    a, b = np.asarray(a), np.asarray(b)
    test.assertEqual(a.dtype, b.dtype)
    test.assertFalse(np.isinf(a).any() or np.isinf(b).any())
    np.testing.assert_array_equal(np.isnan(a), np.isnan(b))
    finite = np.isfinite(a)
    np.testing.assert_array_equal(a[finite].view(np.uint32), b[finite].view(np.uint32))


def accumulations(m):
    return [c for c in m["calls"] if c[0] == "BlockAccumulateExchange"]


class ExchangeWindowCPU(unittest.TestCase):
    def test_constructor_validation(self):
        with mocked_backend() as m:
            with self.assertRaisesRegex(ValueError, "requires observe_raw_exchange"):
                build(m, window=True, raw=False)
            with self.assertRaisesRegex(ValueError, "must be bool"):
                build(m, window=1)

    def test_state_machine_and_launches(self):
        with mocked_backend() as m, ExitStack() as stack:
            stack.enter_context(patch.object(cuda, "memset_d8"))
            stack.enter_context(patch.object(cuda, "memset_d8_async"))
            stack.enter_context(patch.object(cuda, "memcpy_dtoh",
                                             side_effect=lambda out, src: out.fill(0)))
            backend = build(m)
            b = backend._level_boundaries[1]
            self.assertIsNotNone(b.window_integral)
            with self.assertRaises(RuntimeError):
                backend.start_exchange_window()  # not initialized
            arrays = [np.full(level.shape, 1 / 9, np.float32) for level in backend._levels]
            backend.initialize(*arrays)
            with self.assertRaisesRegex(RuntimeError, "no valid exchange window"):
                backend.read_exchange_window(0, level=1)
            backend.step(1)
            self.assertEqual(len(accumulations(m)), 2)  # closed window: interval ledger only
            backend.start_exchange_window()
            with self.assertRaisesRegex(RuntimeError, "at least one completed"):
                backend.read_exchange_window(0, level=1)
            m["calls"].clear()
            backend.step(2)
            calls = accumulations(m)
            self.assertEqual(len(calls), 8)  # two fine hops x two ledgers x two base steps
            self.assertEqual({id(c[1][1]) for c in calls},
                             {id(b.interval_integral), id(b.window_integral)})
            out = backend.read_exchange_window(0, level=1)
            self.assertEqual((out["time_start"], out["time_end"], out["sample_count"]), (1., 3, 4))
            self.assertEqual(backend.read_exchange_window(0, level=1)["time_end"], 3)  # non-destructive
            backend.initialize(*arrays)  # invalidates
            with self.assertRaisesRegex(RuntimeError, "no valid exchange window"):
                backend.read_exchange_window(0, level=1)
            backend.start_exchange_window()
            with patch.object(backend, "_hop", side_effect=RuntimeError("injected")):
                with self.assertRaises(RuntimeError):
                    backend.step(1)
            self.assertFalse(backend._window_valid)
            with self.assertRaises(RuntimeError):
                backend.read_exchange_window(0, level=1)  # failed status

    def test_disabled_default_has_no_window_cost(self):
        with mocked_backend() as m, ExitStack() as stack:
            stack.enter_context(patch.object(cuda, "memset_d8_async"))
            backend = build(m, window=False)
            self.assertIsNone(backend._level_boundaries[1].window_integral)
            backend.initialize(*[np.full(l.shape, 1 / 9, np.float32) for l in backend._levels])
            backend.step(1)
            self.assertEqual(len(accumulations(m)), 2)
            with self.assertRaisesRegex(RuntimeError, "disabled"):
                backend.start_exchange_window()


class ExchangeWindowGPU(unittest.TestCase):
    def test_window_matches_summed_base_reads_and_populations_unchanged(self):
        from test_raw_exchange_gpu import BOUNDS as RB, CFG, circle as rc, jet, state
        from CelerisLab.cuda.context import CudaContext
        with CudaContext(0) as ctx:
            kw = dict(context=ctx, refinements=[RB], level_blocks=[(12, 8), (16, 20)],
                      level_objects=[(), (rc(),)], level_jets=[(), (jet(),)],
                      observe_raw_exchange=True)
            with FixedPatchStepper(CFG, observe_window_exchange=True, **kw) as windowed, \
                 FixedPatchStepper(CFG, **kw) as plain:
                arrays = [state(level.shape) for level in windowed._levels]
                windowed.initialize(*arrays)
                plain.initialize(*arrays)
                windowed.step(2)
                plain.step(2)
                windowed.start_exchange_window()
                b = windowed._level_boundaries[1]
                count = len(b.tables["cell"])
                factors = np.array([.25, .25, .125, .25])  # h*dt, h*h*dt, h*h for level-1 h=.5
                scale = np.zeros(4)
                substeps = 0
                kernel = windowed._levels[1].kernel

                def audited(src, dst, *args, **kw):
                    nonlocal scale, substeps
                    windowed._stream.synchronize()  # BFL and both accumulations are complete
                    raw = np.empty((count, 4), np.float32)
                    cuda.memcpy_dtoh(raw, b.last_raw)
                    scale = scale + (np.abs(raw.astype(np.float64)) * factors).sum(axis=0)
                    substeps += 1
                    kernel(src, dst, *args, **kw)
                windowed._levels[1].kernel = audited
                keys = ("impulse", "angular_impulse", "mass_into_fluid")
                summed = {part: [0., 0., 0., 0.] for part in ("total", "wall", "aperture")}
                steps = 0
                for speed in (0., .002, .002, -.001, 0., 0.):  # amplitude updates stay in-window
                    windowed.set_jet_amplitude(0, speed, level=1)
                    plain.set_jet_amplitude(0, speed, level=1)
                    windowed.step(1)
                    plain.step(1)
                    steps += 1
                    obs = windowed.read_raw_exchange(0, level=1)
                    for part in summed:
                        o = obs[part]
                        summed[part] = [summed[part][0] + o["impulse"][0], summed[part][1] + o["impulse"][1],
                                        summed[part][2] + o["angular_impulse"],
                                        summed[part][3] + o["mass_into_fluid"]]
                    for a, e in zip(windowed._download_raw().values, plain._download_raw().values):
                        same_bits(self, a, e)
                out = windowed.read_exchange_window(0, level=1)
                self.assertEqual((out["time_end"] - out["time_start"], out["sample_count"]),
                                 (steps, 2 * steps))
                self.assertEqual(substeps, 2 * steps)
                bound = 4 * (count * substeps) * EPS * scale + 1e-300
                for part in summed:
                    o = out[part]
                    got = np.array([o["impulse"][0], o["impulse"][1], o["angular_impulse"],
                                    o["mass_into_fluid"]])
                    diff = np.abs(got - np.array(summed[part]))
                    self.assertTrue((diff <= bound).all(), f"{part}: {diff} > {bound}")
                    np.testing.assert_array_equal(o["force"], o["impulse"] / steps)
                again = windowed.read_exchange_window(0, level=1)
                np.testing.assert_array_equal(again["total"]["impulse"], out["total"]["impulse"])
                windowed.initialize(*arrays)
                with self.assertRaises(RuntimeError):
                    windowed.read_exchange_window(0, level=1)


def audit_substeps(stepper, b, count, record):
    """Wrap the level-1 step kernel: after BFL and accumulation, add |substep terms| to record."""
    factors = np.array([.25, .25, .125, .25])
    kernel = stepper._levels[1].kernel

    def audited(src, dst, *args, **kw):
        stepper._stream.synchronize()
        raw = np.empty((count, 4), np.float32)
        cuda.memcpy_dtoh(raw, b.last_raw)
        record["scale"] = record["scale"] + (np.abs(raw.astype(np.float64)) * factors).sum(axis=0)
        record["substeps"] += 1
        kernel(src, dst, *args, **kw)
    stepper._levels[1].kernel = audited


def components(obs):
    o = obs["total"]
    return np.array([o["impulse"][0], o["impulse"][1], o["angular_impulse"], o["mass_into_fluid"]])


class ExchangeWindowShiftedMRTGPU(unittest.TestCase):
    """Shifted MRT, all-blocked, stationary body, NO Jet (scoped separately from SRT/Jet)."""

    def test_window_step0_reset_and_split_step(self):
        from test_raw_exchange_gpu import BOUNDS as RB, circle as rc, state
        from CelerisLab.config import LBMConfig
        from CelerisLab.cuda.context import CudaContext
        cfg = LBMConfig(nx=32, ny=24, viscosity=.1, collision="MRT", ddf_shifting=True)
        with CudaContext(0) as ctx:
            kw = dict(context=ctx, refinements=[RB], level_blocks=[(12, 8), (16, 20)],
                      level_objects=[(), (rc(),)], observe_raw_exchange=True)
            with FixedPatchStepper(cfg, observe_window_exchange=True, **kw) as windowed, \
                 FixedPatchStepper(cfg, observe_window_exchange=True, **kw) as split, \
                 FixedPatchStepper(cfg, **kw) as plain:
                arrays = [state(level.shape) for level in windowed._levels]
                for run in (windowed, split, plain):
                    run.initialize(*arrays)
                    run.step(2)
                    if run is not plain:
                        run.start_exchange_window()
                b = windowed._level_boundaries[1]
                count = len(b.tables["cell"])
                record = dict(scale=np.zeros(4), substeps=0)
                audit_substeps(windowed, b, count, record)
                steps, summed = 5, np.zeros(4)
                for _ in range(steps):
                    windowed.step(1)
                    plain.step(1)
                    summed += components(windowed.read_raw_exchange(0, level=1))
                    for a, e in zip(windowed._download_raw().values, plain._download_raw().values):
                        same_bits(self, a, e)
                split.step(steps)  # one call: same ledger arithmetic as five step(1) calls
                for a, e in zip(windowed._download_raw().values, split._download_raw().values):
                    same_bits(self, a, e)
                rows = [np.empty((count, 4), np.float64) for _ in range(2)]
                cuda.memcpy_dtoh(rows[0], b.window_integral)
                cuda.memcpy_dtoh(rows[1], split._level_boundaries[1].window_integral)
                np.testing.assert_array_equal(rows[0].view(np.uint64), rows[1].view(np.uint64))

                out = windowed.read_exchange_window(0, level=1)
                self.assertEqual(record["substeps"], 2 * steps)
                self.assertEqual((out["time_end"] - out["time_start"], out["sample_count"]),
                                 (steps, 2 * steps))
                bound = 4 * (count * record["substeps"]) * EPS * record["scale"] + 1e-300
                diff = np.abs(components(out) - summed)
                self.assertTrue((diff <= bound).all(), f"{diff} > {bound}")
                np.testing.assert_array_equal(split.read_exchange_window(0, level=1)["total"]["impulse"],
                                              out["total"]["impulse"])

                windowed.step(0)  # no interval: window and its duration unchanged
                again = windowed.read_exchange_window(0, level=1)
                self.assertEqual(again["time_end"], out["time_end"])
                np.testing.assert_array_equal(again["total"]["impulse"], out["total"]["impulse"])

                windowed.start_exchange_window()  # reset after a window
                with self.assertRaisesRegex(RuntimeError, "at least one completed"):
                    windowed.read_exchange_window(0, level=1)
                record.update(scale=np.zeros(4), substeps=0)
                windowed.step(1)
                single = components(windowed.read_raw_exchange(0, level=1))
                reset = windowed.read_exchange_window(0, level=1)
                self.assertEqual(reset["time_end"] - reset["time_start"], 1)
                bound = 4 * (count * record["substeps"]) * EPS * record["scale"] + 1e-300
                self.assertTrue((np.abs(components(reset) - single) <= bound).all())


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