"""Host-compile the actual numerical helper; no CUDA or flow evolution."""
import ctypes
from pathlib import Path
import subprocess
import sys
from tempfile import TemporaryDirectory
import unittest
import numpy as np
ROOT=Path(__file__).resolve().parents[3]
sys.path.insert(0,str(ROOT/'src'))
from CelerisLab.config import LBMConfig, INLET_SCHEME_MAP
cx=np.array([0,1,-1,0,0,1,-1,1,-1]);cy=np.array([0,0,0,1,-1,1,-1,-1,1]);w=np.array([4/9]+[1/9]*4+[1/36]*4)
M=np.array([np.ones(9),cx,cy]);Q=np.array([cx*cx,cy*cy,cx*cy])
TEMP=TemporaryDirectory();libs=[]
for shifted in [0,1]:
    code=f'''
#include <cmath>
#define __device__
#define __forceinline__ inline
#define __fdiv_rn(a,b) ((a)/(b))
#define LATTICE_D2Q9 0
#define LATTICE_MODEL 0
#define NQ 9
#define CS2 (1.0f/3.0f)
#define NY 330
#define U0 .01f
#define W0_VAL (4.0f/9.0f)
#define WS_VAL (1.0f/9.0f)
#define WE_VAL (1.0f/36.0f)
#define USE_DDF_SHIFTING {shifted}
#define INLET_SCHEME 4
#include "operators/macro.cuh"
#include "operators/collision_srt.cuh"
#include "boundary/inlet/common.cuh"
#include "boundary/inlet/experimental_full9.cuh"
extern "C" int apply(float* f,float u,float v,int collision) {{
    apply_experimental_full9_inlet_d2q9(f,u,v);
    int count=0;
    if (collision && inlet_scheme_uses_post_collision_ghost()) {{
        float rho,ux,uy,eq[9]; compute_rho_u(f,rho,ux,uy); compute_feq(rho,ux,uy,eq);
        collide_srt_no_force(f,eq,1.0f/.524f); ++count;
    }}
    return count;
}}
'''
    src=Path(TEMP.name)/f'test{shifted}.cpp';src.write_text(code);lib=src.with_suffix('.so')
    subprocess.run(['g++','-std=c++11','-shared','-fPIC','-O2','-ffp-contract=off','-I',str(ROOT/'src/CelerisLab/lbm/kernels'),str(src),'-o',str(lib)],check=True)
    dll=ctypes.CDLL(str(lib));dll.apply.argtypes=[ctypes.POINTER(ctypes.c_float),ctypes.c_float,ctypes.c_float,ctypes.c_int];dll.apply.restype=ctypes.c_int;libs.append(dll)

def eq(rho,u,v):
    cu=cx*u+cy*v;return w*rho*(1+3*cu+4.5*cu*cu-1.5*(u*u+v*v))
def reference(f,u,v):
    rho=(f[0]+f[3]+f[4]+2*(f[2]+f[6]+f[8]))/(1-u);e=eq(rho,u,v);k=f-e
    pi=np.array([2*(k[2]+k[6]+k[8]),k[3]+k[4]+2*(k[6]+k[8]),2*(k[6]-k[8])])
    stress=w*(4.5*((cx*cx-1/3)*pi[0]+(cy*cy-1/3)*pi[1])+9*cx*cy*pi[2])
    return e,stress,pi

def apply(f,u,v,shifted,collision=0):
    raw=np.asarray(f-shifted*w,dtype=np.float32);count=libs[shifted].apply(raw.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),u,v,collision)
    return raw.astype(float)+shifted*w,count

def test_moments_stress_and_shift():
    rng=np.random.default_rng(92)
    for _ in range(20):
        f=eq(1.02,.003,0)+rng.normal(0,.001,9);e,s,pi=reference(f,.003,.001)
        outputs=[]
        for shifted in [0,1]:
            out,count=apply(f,.003,.001,shifted);outputs.append(out)
            assert count==0
            np.testing.assert_allclose(out,e+s,atol=2e-7,rtol=0)
            np.testing.assert_allclose(M@out,M@e,atol=3e-7,rtol=0)
            np.testing.assert_allclose(Q@(out-e),pi,atol=2e-7,rtol=0)
        np.testing.assert_allclose(*outputs,atol=2e-7,rtol=0)

def test_unknown_inputs_ignored_and_equilibrium():
    for shifted in [0,1]:
        f=eq(1.01,.004,0);out,_=apply(f,.004,0,shifted)
        np.testing.assert_allclose(out,f,atol=2e-7,rtol=0)
        f[[1,5,7]]=np.nan;ignored,_=apply(f,.004,0,shifted)
        assert np.isfinite(ignored).all()
        np.testing.assert_array_equal(out,ignored)

def test_exactly_once_srt():
    f=eq(1,.003,0)+np.array([.001,0,.002,.003,-.001,0,.002,0,-.002]);e,s,_=reference(f,.003,0)
    for shifted in [0,1]:
        out,count=apply(f,.003,0,shifted,1)
        assert count==1
        np.testing.assert_allclose(out,e+(1-1/.524)*s,atol=3e-7,rtol=0)
        assert np.max(abs(out-(e+(1-1/.524)**2*s)))>1e-4
    source=(ROOT/'src/CelerisLab/lbm/kernels/step/one_step_double.cu').read_text()
    part=source.split('#if DIM == 2',2)[-1].split('#elif DIM == 3')[0]
    assert part.index('apply_boundary_pull(f,')<part.index('collide_dispatch(f,')<part.index('stream_pull_store(k,')
    assert part.count('collide_dispatch(f,')==1

def test_config_guard_and_legacy_ids():
    assert {k:INLET_SCHEME_MAP[k] for k in ['zou_he_local','channel_stabilized','equilibrium','regularized']}==dict(zou_he_local=0,channel_stabilized=1,equilibrium=2,regularized=3)
    assert LBMConfig().inlet_scheme=='zou_he_local'
    for shifted in [False,True]:
        c=LBMConfig(inlet_scheme='experimental_regularized_full9',ddf_shifting=shifted);c.validate();assert c.to_macros()['INLET_SCHEME']==4
    for change in [dict(collision='MRT'),dict(collision='TRT'),dict(streaming='esopull'),dict(store_precision='FP16S'),dict(les_enabled=True),dict(y_wall_bc='free_slip'),dict(lattice_model='D3Q19')]:
        with unittest.TestCase().assertRaises(ValueError):LBMConfig(inlet_scheme='experimental_regularized_full9',**change).validate()

if __name__=='__main__':
    result=unittest.TextTestRunner(verbosity=2).run(unittest.TestSuite(unittest.FunctionTestCase(fn) for name,fn in sorted(globals().items()) if name.startswith('test_')))
    TEMP.cleanup();raise SystemExit(not result.wasSuccessful())