"""Real host methods with bounded device-memory doubles; no CUDA context/init.

The memory double checks copy bounds and rejects pageable async transfers.
Kernel math is covered separately by test_numerical_operators; snapshot tests
here check host layout, current-parity arguments and field-only semantics.
"""
from types import SimpleNamespace

import h5py
import numpy as np
import pytest
import pycuda.driver as cuda

from CelerisLab.config import LBMConfig
from CelerisLab.body.manager import ObjectManager
from CelerisLab.body.objects import SimObject
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.ddf_patch import patch_ddf_for_body_sync, upload_patched_ddf
from CelerisLab.common.checkpoint import save_checkpoint, load_checkpoint
from CelerisLab.lbm.field import LBMField
from CelerisLab.lbm.descriptors import FLUID, SOLID, OBSTACLE, SENSOR_FLAG, FRC_REGION
from CelerisLab.lbm.equilibrium import compute_feq_d2q9

W = np.array([4/9] + [1/9]*4 + [1/36]*4, dtype=np.float32)


class Memory:
    def __init__(self, address, size):
        self.address, self.data, self.freed = address, bytearray(size), False

    def __int__(self):
        return self.address

    def free(self):
        assert not self.freed
        self.freed = True


class DeviceDouble:
    def __init__(self):
        self.blocks, self.pinned = [], []
        self.writes, self.syncs = 0, 0

    def alloc(self, size):
        assert size > 0
        block = Memory(1000000*(len(self.blocks)+1), int(size))
        self.blocks.append(block)
        return block

    def locate(self, ptr, size):
        address = int(ptr)
        for block in self.blocks:
            if block.address <= address < block.address+len(block.data):
                offset = address-block.address
                assert not block.freed and offset+size <= len(block.data)
                return block, offset
        raise AssertionError(f"Invalid device address: {address}")

    def htod(self, ptr, array):
        data = array if isinstance(array, bytes) else np.asarray(array).tobytes()
        block, offset = self.locate(ptr, len(data))
        block.data[offset:offset+len(data)] = data
        self.writes += 1

    def dtoh(self, array, ptr):
        block, offset = self.locate(ptr, array.nbytes)
        array[:] = np.frombuffer(block.data, dtype=array.dtype,
                                 count=array.size, offset=offset).reshape(array.shape)

    def pinned_empty(self, shape, dtype):
        array = np.empty(shape, dtype=dtype)
        self.pinned.append(array)
        return array

    def require_pinned(self, array):
        assert any(np.shares_memory(array, item) for item in self.pinned), "pageable async transfer"

    def htod_async(self, ptr, array, stream):
        self.require_pinned(array)
        self.htod(ptr, array)

    def dtoh_async(self, array, ptr, stream):
        self.require_pinned(array)
        self.dtoh(array, ptr)

    def synchronize(self):
        self.syncs += 1


@pytest.fixture
def device(monkeypatch):
    device = DeviceDouble()
    for name, method in {
        "mem_alloc": device.alloc, "memcpy_htod": device.htod,
        "memcpy_dtoh": device.dtoh, "memcpy_htod_async": device.htod_async,
        "memcpy_dtoh_async": device.dtoh_async, "pagelocked_empty": device.pinned_empty,
    }.items():
        monkeypatch.setattr(cuda, name, method)
    monkeypatch.setattr(cuda, "Context", SimpleNamespace(synchronize=device.synchronize))
    return device


class ModuleDouble:
    def __init__(self, device):
        self.device, self.globals, self.calls = device, {}, []
        for name, value in {
            "d_w": W,
            "d_cx": np.array([0,1,-1,0,0,1,-1,1,-1], dtype=np.int32),
            "d_cy": np.array([0,0,0,1,-1,1,-1,-1,1], dtype=np.int32),
            "d_params": np.zeros(5, dtype=np.float32),
        }.items():
            ptr = device.alloc(value.nbytes)
            device.htod(ptr, value)
            self.globals[name] = ptr, value.nbytes

    def get_global(self, name):
        return self.globals[name]

    def get_function(self, name):
        def launch(ddf, scratch, indices, count, step, **kwargs):
            self.calls.append((name, int(step)))
            idx = np.frombuffer(indices.data, dtype=np.uint32)
            raw = np.frombuffer(ddf.data, dtype=np.float32).reshape(9, -1)
            # Deliberately simple invertible parity layout. This is a host
            # protocol oracle, not a substitute for testing CUDA EsoPull math.
            view = raw if int(step) % 2 == 0 else raw[::-1]
            local = np.frombuffer(scratch.data, dtype=np.float32).reshape(int(count), 9)
            if name == "DecodeCellsToPhysical":
                local[:] = view[:, idx].T
            elif name == "EncodePhysicalToCells":
                view[:, idx] = local.T
            else:
                raise AssertionError(name)
        return launch


def make_field(device, **kwargs):
    cfg = LBMConfig(nx=12, ny=10, viscosity=.1, **kwargs)
    cfg.validate()
    field = LBMField(cfg, ModuleDouble(device))
    field.ddf[:] = np.repeat(compute_feq_d2q9(1.1, .03, -.01)[:, None], field.n, axis=1).reshape(-1)
    field.upload_ddf()
    field.upload_flags()
    return field


def obj(x=5, y=5, r=2, **kwargs):
    return SimObject(-1, geometry=CircleGeometry(x, y, r), center=(x,y), radius=r, **kwargs)


def manager(field, objects=()):
    bodies = ObjectManager(field.nx, field.ny, field.nz, field.nq, field.cfg)
    for item in objects:
        bodies.add(item)
    return bodies


@pytest.mark.parametrize("precision", ["FP32", "FP16S"])
@pytest.mark.parametrize("shift", [False, True])
def test_patch_qmajor_storage_and_untouched_bytes(device, precision, shift):
    field = make_field(device, store_precision=precision, ddf_shifting=shift)
    old = field.ddf.copy()
    before = bytes(field.ddf_gpu.data)
    added = np.zeros(field.n, dtype=bool)
    added[[5, 6, field.n-2]] = True
    released = np.zeros_like(added)
    patch_ddf_for_body_sync(field, old, field.flag, field.flag, added, released)
    expected = old.reshape(9, -1).copy()
    expected[:, added] = compute_feq_d2q9(field.cfg.rho, 0, 0)[:, None]
    np.testing.assert_array_equal(field.ddf.reshape(9, -1), expected)
    upload_patched_ddf(field, added, released)
    encoded = (expected - W[:, None]) if shift else expected
    encoded = (encoded * 32768).astype(np.float16) if precision == "FP16S" else encoded
    np.testing.assert_array_equal(np.frombuffer(field.ddf_gpu.data, dtype=field.store_dtype).reshape(9,-1), encoded)
    itemsize = field.store_bytes
    for q in range(9):
        for cell in np.flatnonzero(~added):
            offset = (q*field.n+cell)*itemsize
            assert field.ddf_gpu.data[offset:offset+itemsize] == before[offset:offset+itemsize]
    assert field.ddf_gpu.data == field.temp_gpu.data


def test_released_bfs_uses_qmajor_donors(device):
    field = make_field(device)
    old = field.ddf.copy()
    flags = np.full(field.n, FLUID, dtype=np.uint16)
    released = np.zeros(field.n, dtype=bool)
    released.reshape(field.ny,field.nx)[3:7,3:8] = True
    flags[released] = SOLID
    field.ddf.reshape(9,-1)[:,released] = -99
    patch_ddf_for_body_sync(field, old, flags, field.flag,
                           np.zeros(field.n,dtype=bool), released)
    expected = compute_feq_d2q9(1.1,.03,-.01)
    np.testing.assert_allclose(field.ddf.reshape(9,-1)[:,released],
                               np.repeat(expected[:,None], released.sum(), axis=1), atol=7e-8)


def test_force_survives_other_body_edit_sync_and_preinit_remove(device):
    field = make_field(device)
    bodies = manager(field, [obj(x=3), obj(x=8, is_force_region=True)])
    bodies.set_force_state(1, .004, -.003)
    bodies.set_body_state(0, .2)
    bodies._refresh_action_from_objects()
    np.testing.assert_allclose(bodies.action[6:8], [.004,-.003])
    bodies.remove(0)
    assert bodies.get(0).is_force_region and bodies.get(0).obj_id == 0
    np.testing.assert_allclose(bodies.action[:2], [.004,-.003])
    assert bodies.add(obj(x=3)) == 1
    next_objects = bodies.build_next_objects()
    bodies.commit_pending(next_objects, np.zeros(2, dtype=np.int32))
    bodies._refresh_action_from_objects()
    np.testing.assert_allclose(bodies.action[:2], [.004,-.003])


@pytest.mark.parametrize("reverse", [False, True])
def test_overlay_preserves_occupancy_and_filters_lists(device, reverse):
    field = make_field(device)
    objects = [obj(), obj(r=3,is_sensor=True), obj(r=3,is_force_region=True)]
    if reverse:
        objects.reverse()
    bodies = manager(field, objects)
    flags = bodies.build_flags(field.build_channel_flags())
    center = 5+5*field.nx
    assert flags[center] & SOLID and not flags[center] & FLUID
    assert not flags[center] & (SENSOR_FLAG | FRC_REGION)
    arrays = bodies.build_compact_lists(flags)
    assert all(flags[c] & FLUID and not flags[c] & SOLID for c in arrays[8])
    assert all(flags[c] & FLUID and not flags[c] & SOLID for c in arrays[10])
    sensor = next(o for o in bodies.objects if o.is_sensor)
    assert bodies.sensor_cell_counts[sensor.obj_id] == len(arrays[8])
    assert len(arrays[8]) < len(sensor.get_sensor_list(field.nx, field.ny)[0])


def test_overlapping_force_regions_rejected(device):
    field = make_field(device)
    bodies = manager(field, [obj(is_force_region=True), obj(x=6,is_force_region=True)])
    with pytest.raises(ValueError, match="Overlapping force_region"):
        bodies.build_compact_lists(bodies.build_flags(field.flag))


@pytest.mark.parametrize("second_x,error", [(14., "unique cut-link writers"),
                                           (17., "active fluid receivers")])
def test_passive_overlap_rejected_before_publishing_lists(second_x, error):
    cfg = LBMConfig(nx=32, ny=24)
    cfg.validate()
    bodies = ObjectManager(cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
    for x in (14., second_x):
        bodies.add(obj(x=x, y=12., r=4.3))
    bodies.add(obj(x=4., y=4., r=2., is_sensor=True))
    counts = np.full(bodies.count, 91, np.int32)
    hits = []
    flags = bodies.build_flags(np.full(cfg.nx * cfg.ny, FLUID, np.uint16))
    for domain_flags in (flags, None):
        with pytest.raises(ValueError, match=error):
            bodies.build_compact_lists_for(bodies.objects, domain_flags,
                                          sensor_cell_counts_out=counts, hits_out=hits)
        np.testing.assert_array_equal(counts, 91)
        assert hits == []


@pytest.mark.parametrize("centers", [((10., 12.), (22., 12.)),
                                    ((10.5, 12.), (19.5, 12.))])
def test_separated_passive_bodies_keep_valid_writers(centers):
    # The second case has a subcell physical gap with a resolved fluid column.
    cfg = LBMConfig(nx=32, ny=24)
    cfg.validate()
    bodies = ObjectManager(cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
    for x, y in centers:
        bodies.add(obj(x=x, y=y, r=4.3))
    flags = bodies.build_flags(np.full(cfg.nx * cfg.ny, FLUID, np.uint16))
    indices, directions, _, ids, *_ = bodies.build_compact_lists(flags)
    assert set(ids) == {0, 1}
    assert np.all(flags[indices] & FLUID) and not np.any(flags[indices] & SOLID)
    assert len(np.unique(indices.astype(np.uint64) * cfg.nq + directions)) == len(indices)


@pytest.mark.parametrize("entry", ["sync", "initialize", "runtime_edit"])
@pytest.mark.parametrize("second_x", [5., 6.])
def test_passive_overlap_failure_keeps_live_state(device, entry, second_x):
    from CelerisLab.simulation import Simulation
    field = make_field(device)
    bodies = manager(field, [obj(), obj(x=9., y=3., r=1., is_sensor=True)])
    bodies.sync_to_gpu(field)
    overlapping = obj(x=second_x)
    if entry == "runtime_edit":
        bodies.stage_add(overlapping)
    else:
        bodies.add(overlapping)
    objects = bodies.objects
    flags, curved, counts = field.flag.copy(), field.curved.fluid_idx.copy(), bodies.sensor_cell_counts.copy()
    action = bodies.action.copy()
    buffers = [(block, bytes(block.data), block.freed) for block in device.blocks]
    writes = device.writes
    sim = SimpleNamespace(bodies=bodies, field=field, lbm_cfg=field.cfg,
                          _initialized=entry == "runtime_edit",
                          _assert_allocation_contract=lambda: None,
                          stream=SimpleNamespace(synchronize=device.synchronize))
    sim._require_usable = Simulation._require_usable.__get__(sim)
    with pytest.raises(ValueError, match="active fluid receivers|unique cut-link writers"):
        if entry == "sync":
            bodies.sync_to_gpu(field)
        elif entry == "initialize":
            Simulation.initialize(sim)
        else:
            Simulation.sync_bodies(sim)
    assert bodies.objects == objects
    assert bodies.has_pending_edit() == (entry == "runtime_edit")
    assert sim._initialized == (entry == "runtime_edit")
    assert device.writes == writes and len(device.blocks) == len(buffers)
    assert all(bytes(block.data) == data and block.freed == freed for block, data, freed in buffers)
    np.testing.assert_array_equal(field.flag, flags)
    np.testing.assert_array_equal(field.curved.fluid_idx, curved)
    np.testing.assert_array_equal(bodies.sensor_cell_counts, counts)
    np.testing.assert_array_equal(bodies.action, action)


def test_action_upload_and_scratch_are_pinned(device):
    field = make_field(device, streaming="esopull")
    bodies = manager(field, [obj(is_force_region=True)])
    bodies._allocate_packed_telemetry()
    bodies.set_force_state(0,.001,.002)
    stream = SimpleNamespace(synchronize=device.synchronize)
    bodies._upload_action_async(stream)
    np.testing.assert_array_equal(np.frombuffer(bodies.action_gpu.data,dtype=np.float32), bodies.action)
    scratch, host, indices = field.esopull_sync_decode(np.array([3,8],dtype=np.uint32),1,stream)
    field.esopull_sync_encode(scratch,host,indices,2,1,stream)
    assert scratch.freed and indices.freed and device.syncs >= 2


def test_esopull_snapshot_field_only_across_parity(device):
    field = make_field(device, streaming="esopull")
    original = field.ddf.copy()
    with pytest.raises(ValueError, match="step_id"):
        field.snapshot()
    field.snapshot(step_id=0)
    field.ddf_gpu.data[:] = bytes(len(field.ddf_gpu.data))
    field.restore(step_id=7)
    result = np.frombuffer(field.ddf_gpu.data,dtype=np.float32).reshape(9,-1)[::-1]
    np.testing.assert_array_equal(result, original.reshape(9,-1))
    assert field.module.calls == [("DecodeCellsToPhysical",0),("EncodePhysicalToCells",7)]
    assert field._host_ddf_step is None


def test_snapshot_double_buffer_preserves_field_only_semantics(device):
    field = make_field(device)
    original = field.ddf.copy()
    field.snapshot()
    field.ddf *= 2
    field.upload_ddf()
    field.restore()
    field.download_ddf(force=True)
    np.testing.assert_array_equal(field.ddf,original)


def test_checkpoint_restores_action_cache_and_observation(device,tmp_path):
    field = make_field(device)
    bodies = manager(field,[obj(x=3,r=1),obj(x=8,r=1,is_force_region=True)])
    bodies.set_body_state(0,.125)
    bodies.set_force_state(1,.001,-.002)
    bodies.sync_to_gpu(field)
    field.update_params(fx=.0002)
    bodies.obs_pinned[:] = np.arange(bodies.obs_total_floats)
    device.htod(bodies.obs_gpu,bodies.obs_pinned)
    bodies._obs_accum_steps = 4
    stepper = SimpleNamespace(step_count=7,_step_count=7)
    path = save_checkpoint(field,stepper,field.cfg,bodies,str(tmp_path/'state.h5'))
    expected = bytes(field.ddf_gpu.data)
    bodies.set_body_state(0,0)
    bodies.set_force_state(1,0,0)
    field.ddf[:] = 99
    field._host_ddf_step = 7
    field.ddf_gpu.data[:] = bytes(len(field.ddf_gpu.data))
    load_checkpoint(path,field,stepper,field.cfg,bodies)
    assert bytes(field.ddf_gpu.data) == expected
    assert field._host_ddf_step is None
    field.download_ddf(step_id=7)
    assert field.ddf.max() < 1
    np.testing.assert_allclose(bodies.action[[5,6,7]],[.125,.001,-.002])
    np.testing.assert_array_equal(np.frombuffer(bodies.action_gpu.data,dtype=np.float32),bodies.action)
    assert bodies._obs_accum_steps == 4
    np.testing.assert_array_equal(bodies.obs_pinned,np.arange(bodies.obs_total_floats))
    assert field._runtime_params['fx'] == .0002


@pytest.mark.parametrize("periodic_x", [False, True])
@pytest.mark.parametrize("with_objects", [False, True])
def test_checkpoint_restores_actual_domain(device, tmp_path, periodic_x, with_objects):
    source = make_field(device)
    objects = [obj(x=2, r=1), obj(x=10, r=2, is_sensor=True)] if with_objects else []
    bodies = manager(source, objects)
    base = source.build_channel_flags()
    if periodic_x:
        base.reshape(source.ny, source.nx)[1:-1, [0, -1]] = FLUID
    source.flag = bodies.build_flags(base)
    source.upload_flags()
    bodies.sync_to_gpu(source, rebuild_flags=False)
    saved_flags = source.flag.copy()
    # The device, not a possibly stale host mirror, is authoritative at save.
    source.flag = source.build_channel_flags()
    path = save_checkpoint(source, SimpleNamespace(step_count=3), source.cfg,
                           bodies, str(tmp_path / 'domain.h5'))
    target = make_field(device)
    target_bodies = manager(target, objects)
    stepper = SimpleNamespace(_step_count=0)
    load_checkpoint(path, target, stepper, target.cfg, target_bodies)
    assert stepper._step_count == 3
    np.testing.assert_array_equal(target.flag, saved_flags)
    np.testing.assert_array_equal(np.frombuffer(target.flag_gpu.data, np.uint16), saved_flags)
    assert target.ddf_gpu.data == source.ddf_gpu.data


@pytest.mark.parametrize("periodic_x", [False, True])
@pytest.mark.parametrize("change", ["geometry", "role", "ddf_shape"])
def test_checkpoint_rejects_before_device_writes(device,tmp_path,change,periodic_x):
    field = make_field(device)
    bodies = manager(field,[obj()])
    bodies.sync_to_gpu(field)
    if periodic_x:
        field.flag.reshape(field.ny, field.nx)[1:-1, [0, -1]] = FLUID
        field.upload_flags()
    path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'state.h5'))
    if change == 'geometry':
        bodies.get(0).geometry = CircleGeometry(6,5,2)
    elif change == 'role':
        bodies.get(0)._is_sensor = True
    else:
        with h5py.File(path,'a') as hf:
            del hf['ddf']
            hf.create_dataset('ddf',data=np.zeros(2,dtype=np.float32))
    writes = device.writes
    with pytest.raises(ValueError):
        load_checkpoint(path,field,SimpleNamespace(_step_count=3),field.cfg,bodies)
    assert device.writes == writes


def test_checkpoint_parameter_study_keeps_new_viscosity(device,tmp_path):
    field = make_field(device)
    bodies = manager(field)
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'state.h5'))
    field.cfg.viscosity = .2
    field.cfg.validate()
    load_checkpoint(path,field,SimpleNamespace(_step_count=0),field.cfg,bodies)
    assert field._runtime_params['omega'] == field.cfg.omega


@pytest.mark.parametrize('version', [1, 2])
def test_legacy_checkpoint_default_domain(device, tmp_path, version):
    field = make_field(device)
    bodies = manager(field, [obj()])
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg,
                           bodies, str(tmp_path / 'legacy-domain.h5'))
    with h5py.File(path, 'a') as hf:
        hf.attrs['ckpt_version'] = version
        del hf['domain_flag']
    load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies)
    # Legacy files have no authoritative base-domain record. Keep their check.
    with h5py.File(path, 'a') as hf:
        hf['flag'][field.nx] = FLUID
    writes = device.writes
    with pytest.raises(ValueError, match='geometry'):
        load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies)
    assert device.writes == writes


@pytest.mark.parametrize('damage', ['missing', 'shape', 'dtype', 'boundary', 'hidden', 'annotation'])
def test_checkpoint_rejects_invalid_domain_before_writes(device, tmp_path, damage):
    field = make_field(device)
    bodies = manager(field, [obj()])
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg,
                           bodies, str(tmp_path / 'invalid-domain.h5'))
    with h5py.File(path, 'a') as hf:
        if damage in ('missing', 'shape', 'dtype'):
            del hf['domain_flag']
            if damage != 'missing':
                hf.create_dataset('domain_flag', data=np.ones(
                    2 if damage == 'shape' else field.n,
                    dtype=np.uint16 if damage == 'shape' else np.float32))
        elif damage == 'boundary':
            hf['domain_flag'][field.nx] = FLUID
        elif damage == 'hidden':
            hf['domain_flag'][5 + 5 * field.nx] = SOLID
        else:
            # Mutate both records: overlay idempotence alone would accept this.
            for name in ('domain_flag', 'flag'):
                hf[name][field.nx + 1] = FLUID | SENSOR_FLAG
    before = bytes(field.ddf_gpu.data), field.flag.copy()
    writes = device.writes
    with pytest.raises(ValueError):
        load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies)
    assert device.writes == writes
    assert bytes(field.ddf_gpu.data) == before[0]
    np.testing.assert_array_equal(field.flag, before[1])


@pytest.mark.parametrize('extra', [OBSTACLE | SOLID, SENSOR_FLAG, FRC_REGION])
def test_checkpoint_save_rejects_unregistered_overlay(device, tmp_path, extra):
    field = make_field(device)
    bodies = manager(field)
    field.flag[field.nx + 1] |= extra
    field.upload_flags()
    path = tmp_path / 'unchanged.h5'
    path.write_bytes(b'existing checkpoint')
    with pytest.raises(ValueError):
        save_checkpoint(field, SimpleNamespace(step_count=0), field.cfg, bodies, str(path))
    assert path.read_bytes() == b'existing checkpoint'


def test_checkpoint_domain_handles_plain_solid_mask(device, tmp_path):
    class PlainSolidCircle(CircleGeometry):
        def build_flag_mask(self, nx, ny):
            mask = super().build_flag_mask(nx, ny)
            mask[(mask & SOLID) != 0] = SOLID
            return mask

    field = make_field(device)
    body = obj()
    body.geometry = PlainSolidCircle(5, 5, 2)
    bodies = manager(field, [body])
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field, SimpleNamespace(step_count=3), field.cfg,
                           bodies, str(tmp_path / 'plain-solid.h5'))
    with h5py.File(path) as hf:
        assert hf['domain_flag'][5 + 5 * field.nx] == FLUID
    load_checkpoint(path, field, SimpleNamespace(_step_count=0), field.cfg, bodies)


@pytest.mark.parametrize("kind",['circle','sensor','force_region'])
def test_legacy_checkpoint_supported_scope(device,tmp_path,kind):
    field = make_field(device)
    bodies = manager(field,[obj(is_sensor=kind=='sensor',is_force_region=kind=='force_region')])
    bodies.set_body_state(0,.12)
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'legacy.h5'))
    with h5py.File(path,'a') as hf:
        hf.attrs['ckpt_version'] = 1
        del hf.attrs['geometry_json']
        del hf.attrs['runtime_params_json']
        del hf['obs']
    if kind == 'force_region':
        with pytest.raises(ValueError,match='Legacy v1.*force_region'):
            load_checkpoint(path,field,SimpleNamespace(_step_count=0),field.cfg,bodies)
    else:
        bodies.set_body_state(0,0)
        load_checkpoint(path,field,SimpleNamespace(_step_count=0),field.cfg,bodies)
        assert bodies.action[5] == np.float32(.12)
        assert bodies._obs_accum_steps == 0


def test_compact_upload_stream_uses_pinned_staging(device):
    field = make_field(device)
    bodies = manager(field,[obj(),obj(x=9,r=1,is_sensor=True)])
    bodies.sync_to_gpu(field)
    stream = SimpleNamespace(synchronize=device.synchronize)
    field.upload_compact_lists(stream=stream)
    assert device.syncs >= 2
    np.testing.assert_array_equal(np.frombuffer(field.curved.fluid_idx_gpu.data,dtype=np.uint32),
                                  field.curved.fluid_idx)
    np.testing.assert_array_equal(np.frombuffer(field.sensors.cells_gpu.data,dtype=np.uint32),
                                  field.sensors.cells)


@pytest.mark.parametrize('attribute,value',[
    ('runtime_params_json','{"fx": "not-a-number"}'),('obs_accum_steps','broken'),
])
def test_checkpoint_bad_metadata_rejected_before_writes(device,tmp_path,attribute,value):
    field = make_field(device)
    bodies = manager(field)
    bodies.sync_to_gpu(field)
    path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'state.h5'))
    with h5py.File(path,'a') as hf:
        hf.attrs[attribute] = value
    writes = device.writes
    with pytest.raises(ValueError):
        load_checkpoint(path,field,SimpleNamespace(_step_count=3),field.cfg,bodies)
    assert device.writes == writes


def test_checkpoint_preserves_pending_vs_applied_action(device,tmp_path):
    field = make_field(device)
    bodies = manager(field,[obj(is_force_region=True)])
    bodies.set_force_state(0,.001,.002)
    bodies.sync_to_gpu(field)
    applied = bodies.action.copy()
    bodies.set_force_state(0,.003,.004)  # Next upload target; device is unchanged.
    pending = bodies.action.copy()
    path = save_checkpoint(field,SimpleNamespace(step_count=3),field.cfg,bodies,str(tmp_path/'state.h5'))
    bodies.set_force_state(0,0,0)
    load_checkpoint(path,field,SimpleNamespace(_step_count=0),field.cfg,bodies)
    np.testing.assert_array_equal(bodies.action,pending)
    np.testing.assert_array_equal(np.frombuffer(bodies.action_gpu.data,dtype=np.float32),applied)