"""Bounded GPU invariants. No high-Re/benchmark-performance claim.

Compare spatial postcollision populations, not DB macro fields against EsoPull's
next-pull diagnostics. Closed-channel evolution starts at completed count 1 with
an explicitly identical spatial state. The distinct t=0 init contract is tested
separately through the actual initialization and wall staging kernels.
"""
import numpy as np
import pytest
import pycuda.driver as cuda

from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID, BC_WALL
from ._cases import make_sim

CX = np.asarray(D2Q9_EX)
CY = np.asarray(D2Q9_EY)
W = np.asarray([4/9]+[1/9]*4+[1/36]*4, dtype=np.float32)


def equilibrium(rho, ux, uy):
    cu = CX[:, None, None]*ux + CY[:, None, None]*uy
    return (W[:, None, None]*rho*(1+3*cu+4.5*cu*cu-1.5*(ux*ux+uy*uy))).astype(np.float32)


def postcollision_to_backing(values, completed_count):
    """Independent spatial encoder for the previous advancing-store contract."""
    previous = completed_count-1
    result = np.empty_like(values)
    result[0] = values[0]
    for i in range(1,9,2):
        result[i+1 if previous % 2 else i] = np.roll(values[i], (CY[i],CX[i]), (0,1))
        result[i if previous % 2 else i+1] = values[i+1]
    return result


def backing_to_postcollision(values, completed_count):
    previous = completed_count-1
    result = np.empty_like(values)
    result[0] = values[0]
    for i in range(1,9,2):
        result[i] = np.roll(values[i+1 if previous % 2 else i], (-CY[i],-CX[i]), (0,1))
        result[i+1] = values[i if previous % 2 else i+1]
    return result


def install_spatial_state(sim, physical, count=1):
    field = sim.field
    data = physical-W[:,None,None] if sim.lbm_cfg.ddf_shifting else physical.copy()
    if sim.lbm_cfg.streaming == "esopull":
        data = postcollision_to_backing(data,count)
    with sim.ctx.activate():
        cuda.memcpy_htod(field.ddf_gpu,np.ascontiguousarray(data,dtype=np.float32))
        cuda.memcpy_htod(field.temp_gpu,np.ascontiguousarray(data,dtype=np.float32))
    sim.stepper._step_count=count
    field.invalidate_host_ddf_cache()
    # x is periodic; no ghost inlet/outlet donor algorithm is under test.
    field.flag[:]=FLUID
    field.flag[:field.nx]=SOLID|BC_WALL
    field.flag[-field.nx:]=SOLID|BC_WALL
    field.upload_flags()


def spatial_state(sim):
    field=sim.field
    data=np.empty((9,field.ny,field.nx),dtype=np.float32)
    with sim.ctx.activate():
        cuda.memcpy_dtoh(data,field.ddf_gpu)
    if sim.lbm_cfg.streaming == "esopull":
        data=backing_to_postcollision(data,sim.stepper.step_count)
    if sim.lbm_cfg.ddf_shifting:
        data=data+W[:,None,None]
    return data


@pytest.mark.parametrize("collision,shifted,wall", [
    ("SRT",False,"bounce_back"), ("TRT",True,"bounce_back"),
    ("MRT",False,"free_slip"), ("MRT",True,"bounce_back"),
])
def test_nonuniform_closed_channel_matches_between_layouts(collision,shifted,wall):
    traces=[]
    for streaming in ["double_buffer","esopull"]:
        sim=make_sim(nx=32,ny=16,streaming=streaming,shifted=shifted,
                     collision=collision,y_wall=wall)
        try:
            sim.initialize()
            y,x=np.mgrid[:16,:32]
            physical=equilibrium(1+.001*np.cos(.3*x+.2*y),
                                 .01*np.sin(.2*x+.3*y),.005*np.cos(.3*x-.2*y))
            install_spatial_state(sim,physical,count=1)
            trace=[]
            for _ in range(4):  # first launch, both parities, both adjacent rows
                sim.run(1)
                trace.append(spatial_state(sim)[:,1:-1,:])
            traces.append(trace)
        finally:
            sim.close()
    # FP32 local arithmetic/roundoff; old donor-stage errors are much larger.
    np.testing.assert_allclose(traces[0],traces[1],atol=2e-6,rtol=0)


@pytest.mark.parametrize("shifted",[False,True])
@pytest.mark.parametrize("wall",["bounce_back","free_slip"])
def test_actual_initialization_wall_staging(shifted,wall):
    sim=make_sim(nx=16,ny=32,streaming="esopull",shifted=shifted,y_wall=wall)
    try:
        sim.initialize()
        field=sim.field
        cells=np.arange(field.n,dtype=np.uint32)
        scratch, physical, indices=field.esopull_sync_decode(cells,0)
        physical=physical.reshape(field.n,9).T.reshape(9,32,16).copy()
        with sim.ctx.activate():
            scratch.free();indices.free()
            sim.stepper.prepare_boundaries_fn(field.ddf_gpu,sim.stepper.boundary_values_gpu,
                np.uint64(0),block=sim.stepper.block,grid=sim.stepper.boundary_grid)
            prepared=np.empty(6*16+2*32*9,dtype=np.float32)
            cuda.memcpy_dtoh(prepared,sim.stepper.boundary_values_gpu)
            staged=prepared[:6*16].reshape(2,16,3)
        for side,row in [(0,1),(1,30)]:
            donors=[4,6,7] if side==0 else [3,5,8]
            if wall=="free_slip": donors[1:]=donors[1:][::-1]
            expected=physical[donors,row,:].T
            if shifted: expected=expected-W[donors]
            np.testing.assert_allclose(staged[side],expected,atol=1e-7,rtol=0)
        donors=prepared[6*16:].reshape(2,32,9)
        for side,x in [(0,1),(1,14)]:
            expected=physical[:,:,x].T
            if shifted: expected=expected-W
            np.testing.assert_allclose(donors[side],expected,atol=1e-7,rtol=0)
        sim.run(1)  # actual first step consumes the staged values without a race
        assert np.isfinite(spatial_state(sim)).all()
    finally:
        sim.close()


@pytest.mark.parametrize("collision",["SRT","TRT","MRT"])
@pytest.mark.parametrize("streaming",["double_buffer","esopull"])
def test_uniform_global_force_momentum(collision,streaming):
    sim=make_sim(nx=32,ny=16,collision=collision,streaming=streaming,
                 shifted=collision!="SRT")
    try:
        sim.initialize()
        initial=equilibrium(np.full((16,32),1.17),np.full((16,32),.03),np.full((16,32),-.02))
        install_spatial_state(sim,initial)
        sim.field.update_params(fx=.001,fy=-.0007)
        sim.run(1)
        # Stay outside the wall influence region. Global force is uniform.
        delta=(spatial_state(sim)-initial)[:,3:-3,3:-3]
        np.testing.assert_allclose(delta.sum(axis=0),0,atol=5e-7,rtol=0)
        np.testing.assert_allclose(np.einsum('q,qyx->yx',CX,delta),.001,atol=2e-7,rtol=0)
        np.testing.assert_allclose(np.einsum('q,qyx->yx',CY,delta),-.0007,atol=2e-7,rtol=0)
    finally:
        sim.close()


@pytest.mark.parametrize("streaming",["double_buffer","esopull"])
@pytest.mark.parametrize("shifted",[False,True])
def test_force_region_impulse_budget(streaming,shifted):
    sim=make_sim(nx=32,ny=16,streaming=streaming,shifted=shifted,collision="MRT")
    try:
        sim.add_body("force_region",center=(16,8),radius=2)
        sim.initialize()
        initial=equilibrium(np.full((16,32),1.17),np.full((16,32),.03),np.full((16,32),-.02))
        install_spatial_state(sim,initial)
        sim.set_force(0,fx=.001,fy=-.0007)
        sim.run(1)
        # The region postpass uses the public layout-specific current-state
        # convention. Sum global momentum so transport cannot move the target
        # footprint outside a diagnostic slice. Subtract an identical zero-force
        # run to cancel wall impulse and baseline FP32 collision roundoff.
        forced=spatial_state(sim)
        install_spatial_state(sim,initial)
        sim.set_force(0,fx=0,fy=0)
        sim.run(1)
        delta=forced-spatial_state(sim)
        count=sim.field.n_force_region
        assert count>0
        # Budget scales with touched cells and FP32 subtraction error.
        tol=3e-7*count
        assert abs(delta.sum(dtype=np.float64))<tol
        assert abs(np.einsum('q,qyx->',CX,delta)-count*.001)<tol
        assert abs(np.einsum('q,qyx->',CY,delta)+count*.0007)<tol
    finally:
        sim.close()


@pytest.mark.parametrize("omega",[0.,.002])
@pytest.mark.parametrize("periodic",[True,False])
def test_curved_layout_equivalence_isolates_open_boundary(omega,periodic):
    """Both periodic and open cases must agree at the same spatial phase."""
    traces=[]
    for streaming in ["double_buffer","esopull"]:
        sim=make_sim(nx=64,ny=32,streaming=streaming,collision="MRT",y_wall="free_slip")
        try:
            # Match the boundary family of the rotating-cylinder K2 benchmark.
            sim.lbm_cfg.inlet_scheme="regularized"
            sim.lbm_cfg.inlet_profile="uniform"
            sim.lbm_cfg.inlet_regularized_neq_damp=1.
            sim.add_body("circle",center=(32,16),radius=4)
            sim.recompile()
            sim.initialize()
            flags=sim.field.flag.copy()
            y,x=np.mgrid[:32,:64]
            initial=equilibrium(1+.0001*np.cos(.2*x+.1*y),
                .02+.001*np.sin(.3*x+.2*y),.001*np.cos(.2*x-.1*y))
            install_spatial_state(sim,initial,count=1)
            sim.field.flag[:]=flags  # preserve obstacle and curved flags
            if periodic:
                grid=sim.field.flag.reshape(32,64)
                grid[1:-1,0]=FLUID;grid[1:-1,-1]=FLUID
            sim.field.upload_flags()
            sim.set_body(0,omega=omega)
            trace=[]
            for step in range(20):
                sim.run(1)
                if step in [0,1,4,19]: trace.append(spatial_state(sim))
            traces.append(trace)
        finally:
            sim.close()
    np.testing.assert_allclose(traces[0],traces[1],atol=2e-6,rtol=0)