"""Actual CUDA operators compiled with a CPU shim; no GPU/import side effects."""
from pathlib import Path
import shutil
import subprocess
import numpy as np
import pytest

ROOT = Path(__file__).resolve().parents[2]
KERNELS = ROOT / "src/CelerisLab/lbm/kernels"
CPP = Path(__file__).with_name("cpp")


def compile_run(tmp_path, source, defines):
    if not shutil.which("g++"):
        pytest.skip("g++ required for actual-header CPU regression")
    exe = tmp_path / "operator"
    command = ["g++", "-std=c++17", "-O0", f"-I{KERNELS}"]
    command += [f"-D{k}={v}" for k, v in defines.items()]
    subprocess.run(command + [str(CPP / source), "-o", str(exe)], check=True, capture_output=True, text=True)
    return subprocess.check_output([str(exe)], text=True).splitlines()


@pytest.mark.parametrize("dim", [2, 3])
@pytest.mark.parametrize("shift", [0, 1])
@pytest.mark.parametrize("model", [0, 1, 2])
def test_collision_moments_and_inlet_closure(tmp_path, dim, shift, model):
    lines = compile_run(tmp_path, "collision_regression.cpp", dict(
        DIM=dim, LATTICE_MODEL=dim-2, USE_DDF_SHIFTING=shift,
        COLLISION_MODEL=model, NY=8, U0=".03f"))
    cx = np.array([0,1,-1,0,0,1,-1,1,-1] if dim == 2 else
                  [0,1,-1,0,0,0,0,1,-1,1,-1,0,0,1,-1,1,-1,0,0])
    cy = np.array([0,0,0,1,-1,1,-1,-1,1] if dim == 2 else
                  [0,0,0,1,-1,0,0,1,-1,0,0,1,-1,-1,1,0,0,1,-1])
    cz = np.zeros(9) if dim == 2 else np.array([0,0,0,0,0,1,-1,0,0,1,-1,1,-1,0,0,-1,1,-1,1])
    weights = np.array([4/9]+[1/9]*4+[1/36]*4 if dim == 2 else [1/3]+[1/18]*6+[1/36]*12)
    # Independent D2 moment matrix (paired direction ordering).
    M = np.array([[1]*9,[-4,-1,-1,-1,-1,2,2,2,2],
                  [4,-2,-2,-2,-2,1,1,1,1],[0,1,-1,0,0,1,-1,1,-1],
                  [0,-2,2,0,0,1,-1,1,-1],[0,0,0,1,-1,1,-1,-1,1],
                  [0,0,0,-2,2,1,-1,-1,1],[0,1,1,-1,-1,0,0,0,0],
                  [0,0,0,0,0,1,1,-1,-1]], dtype=float)
    for a, omega in enumerate([.7,1.25,1.8]):
        before, eq, source, after = [np.fromstring(s, sep=" ") for s in lines[4*a:4*a+4]]
        delta = after-before
        assert abs(delta.sum()) < 5e-7
        np.testing.assert_allclose([delta@cx,delta@cy], [.001,-.0007], atol=1.5e-7, rtol=0)
        if dim == 3:
            assert abs(delta@cz-.0003) < 1.5e-7
        if model == 0:
            expected=before-omega*(before-eq)+(1-omega/2)*source
        elif model == 1:
            opp=np.array([0]+[i+1 if i%2 else i-1 for i in range(1,len(before))])
            wm=1/(.1875/(1/omega-.5)+.5)
            delta_s=(before-eq+(before-eq)[opp])/2
            delta_a=(before-eq-(before-eq)[opp])/2
            expected=before-omega*delta_s-wm*delta_a+(1-omega/2)*(source+source[opp])/2+(1-wm/2)*(source-source[opp])/2
        elif dim == 2:
            rates=np.array([0,1.2,1.2,0,1.2,0,1.2,omega,omega])
            expected=before+np.linalg.solve(M,-rates*(M@(before-eq))+(1-rates/2)*(M@source))
        else:
            expected=None  # D3 projection MRT: conservation/equilibrium only.
        if expected is not None:
            np.testing.assert_allclose(after,expected,atol=2e-7,rtol=0)
    assert float(lines[12].split()[1]) == pytest.approx(1.17, abs=3e-7)
    velocities=np.stack([cx,cy,cz],axis=1)
    u=np.array([.03,-.02,.01 if dim==3 else 0])
    cu=velocities@u
    eq=weights*1.17*(1+3*cu+4.5*cu**2-1.5*(u@u))-shift*weights
    np.testing.assert_allclose(np.fromstring(lines[13],sep=" "),eq,atol=2e-7,rtol=0)


def test_shifted_mrt_weak_force_full_vector(tmp_path):
    """Small populations must retain their precision through MRT relaxation."""
    lines = compile_run(tmp_path, "collision_regression.cpp", dict(
        DIM=2, LATTICE_MODEL=0, USE_DDF_SHIFTING=1, COLLISION_MODEL=2,
        MRT_WEAK_FORCE_CHECK=1, NY=8, U0=".03f"))
    M = np.array([[1]*9,[-4,-1,-1,-1,-1,2,2,2,2],
                  [4,-2,-2,-2,-2,1,1,1,1],[0,1,-1,0,0,1,-1,1,-1],
                  [0,-2,2,0,0,1,-1,1,-1],[0,0,0,1,-1,1,-1,-1,1],
                  [0,0,0,-2,2,1,-1,-1,1],[0,1,1,-1,-1,0,0,0,0],
                  [0,0,0,0,0,1,1,-1,-1]], dtype=float)
    weights = np.array([4/9]+[1/9]*4+[1/36]*4)
    np.testing.assert_allclose(M @ weights, [1,-2,1,0,0,0,0,0,0], atol=3e-16)
    for start in range(0, len(lines), 4):
        # Nine printed digits round-trip FP32; recover those exact binary inputs
        # before doing FP64 arithmetic (especially rho-1 near unit density).
        omega, rho, ux, uy, fx, fy = np.fromstring(lines[start], sep=" ", dtype=np.float32).astype(float)
        before, source, after = [np.fromstring(s, sep=" ", dtype=np.float32).astype(float)
                                  for s in lines[start+1:start+4]]
        # Independent LL equilibrium moments, with the exact background removed.
        drho, u2 = rho-1, ux*ux+uy*uy
        meq = np.array([drho, -2*drho+3*rho*u2, drho-3*rho*u2,
                        rho*ux, -rho*ux, rho*uy, -rho*uy,
                        rho*(ux*ux-uy*uy), rho*ux*uy])
        rates = np.array([0, np.float32(1.2), np.float32(1.2), 0,
                          np.float32(1.2), 0, np.float32(1.2), omega, omega])
        expected = before + np.linalg.solve(
            M, rates*(meq-M@before)+(1-rates/2)*(M@source))
        # FP32 roundoff is relative to the stored small populations, not to w.
        budget = 8*np.finfo(np.float32).eps * max(abs(before).max(), abs(source).max()) + 1e-13
        np.testing.assert_allclose(after, expected, atol=budget, rtol=0)
        np.testing.assert_allclose((after-before) @ M[[0,3,5]].T,
                                   [0,fx,fy], atol=4*budget, rtol=0)
        if not before.any() and fx == 0 and fy == 0:
            np.testing.assert_array_equal(after, np.zeros(9))


@pytest.mark.parametrize("shift", [0, 1])
@pytest.mark.parametrize("wall", [0, 1])
@pytest.mark.parametrize("model", [0, 1, 2])
def test_actual_semantic_and_auxiliary_kernels(tmp_path, shift, wall, model):
    lines = compile_run(tmp_path, "semantic_regression.cpp", dict(USE_DDF_SHIFTING=shift,Y_WALL_BC=wall,COLLISION_MODEL=model))
    assert "passed" in lines[-1]


@pytest.mark.parametrize("model", [0, 1, 2])
@pytest.mark.parametrize("shift", [0, 1])
@pytest.mark.parametrize("storage", [0, 1])
def test_d2_entrypoint_compiles_in_isolation(tmp_path, model, shift, storage):
    """Both streaming entrypoints compile; generated repository headers untouched."""
    import re
    nvcc = shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc"
    if not Path(nvcc).is_file():
        pytest.skip("CUDA toolkit required for compile-only regression")
    copy = tmp_path / "kernels"
    shutil.copytree(KERNELS, copy)
    for file, defines in [
        ("config_grid.h", {"LATTICE_MODEL": 0, "NX": 16, "NY": 8, "NZ": 1}),
        ("config_method.h", {"COLLISION_MODEL": model, "USE_DDF_SHIFTING": shift,
                             "STORE_PRECISION": storage, "Y_WALL_BC": shift})
    ]:
        path=copy / "config" / file
        source=path.read_text()
        for key,value in defines.items():
            source=re.sub(r"(#define\s+"+key+r"\s+)\S+",lambda match: match[1]+str(value),source)
        path.write_text(source)
    result=subprocess.run([nvcc,"-ptx","-arch=sm_70","-std=c++17",
                           str(copy/"kernel_v2.cu"),"-o",str(tmp_path/"kernel.ptx")],
                          capture_output=True,text=True)
    assert result.returncode == 0, result.stderr