"""GPU ownership, isolated builds and immutable allocation contracts."""
from pathlib import Path

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

import CelerisLab
from CelerisLab.cuda import compiler_v2
from ._cases import make_sim


def test_source_provenance():
    root = Path(__file__).resolve().parents[2] / "src/CelerisLab"
    assert Path(CelerisLab.__file__).resolve().parent == root
    assert Path(compiler_v2.kernel_path("kernel_v2.cu")).resolve() == root / "lbm/kernels/kernel_v2.cu"


def test_two_contexts_interleaved_and_non_lifo_close():
    before = cuda.Context.get_current()
    first = make_sim(nx=32, ny=16)
    second = make_sim(nx=48, ny=24, streaming="esopull", shifted=True)
    try:
        assert first._build_dir != second._build_dir
        assert cuda.Context.get_current() == before
        first.initialize()
        second.initialize()
        for _ in range(3):
            first.run(1)
            second.run(1)
        assert first.get_macroscopic()["rho"].shape == (16, 32)
        assert second.get_macroscopic()["rho"].shape == (24, 48)
        first.close()
        second.run(1)
        assert np.isfinite(second.get_macroscopic()["rho"]).all()
        assert cuda.Context.get_current() == before
    finally:
        second.close()
        first.close()


def test_structural_change_rejected_before_running():
    with make_sim(nx=32, ny=16) as sim:
        sim.initialize()
        before = sim.get_ddf().copy()
        sim.lbm_cfg.nx = 64
        with pytest.raises(ValueError, match="allocation|structural"):
            sim.run(1)
        with pytest.raises(ValueError, match="allocation|structural"):
            sim._recompile()
        sim.lbm_cfg.nx = 32
        np.testing.assert_array_equal(sim.get_ddf(), before)
        assert sim.stepper.step_count == 0
        sim.run(1)


def test_viscosity_recompile_updates_device_omega():
    import struct
    with make_sim(nx=32, ny=16) as sim:
        sim.lbm_cfg.viscosity = 0.08
        sim.recompile()
        assert sim.stepper.step_count == 0
        sim.initialize()
        with sim.ctx.activate():
            ptr, _ = sim._module.get_global("d_params")
            buf = np.empty(20, dtype=np.uint8)
            cuda.memcpy_dtoh(buf, ptr)
        omega = struct.unpack("=ffffI", buf.tobytes())[0]
        assert omega == pytest.approx(sim.lbm_cfg.omega, rel=1e-7)
        sim.run(1)


@pytest.mark.parametrize("streaming",["double_buffer","esopull"])
@pytest.mark.parametrize("prior_steps",[3,4])
def test_reinitialize_matches_fresh_state_and_observation_window(streaming,prior_steps):
    with make_sim(nx=32,ny=16,streaming=streaming,shifted=True) as fresh, \
         make_sim(nx=32,ny=16,streaming=streaming,shifted=True) as reused:
        for sim in (fresh,reused):
            sim.add_body("sensor",center=(16,8),radius=2)
            sim.initialize()
        reused.run(prior_steps)
        assert reused.stepper.step_count==prior_steps
        assert reused.bodies._obs_accum_steps==prior_steps
        reused.initialize()
        assert reused.stepper.step_count==0
        assert reused.bodies._obs_accum_steps==0
        np.testing.assert_array_equal(reused.bodies.obs_pinned,0)
        with reused.ctx.activate():
            obs=np.empty(reused.bodies.obs_total_floats,np.float32)
            cuda.memcpy_dtoh(obs,reused.bodies.obs_gpu)
        np.testing.assert_array_equal(obs,0)
        np.testing.assert_array_equal(reused.get_ddf(),fresh.get_ddf())
        for key,value in fresh.get_macroscopic().items():
            np.testing.assert_array_equal(reused.get_macroscopic()[key],value)
        # Equal next step proves restored parity, not only equal initialization bytes.
        fresh.run(1);reused.run(1)
        np.testing.assert_array_equal(reused.get_ddf(),fresh.get_ddf())