# Равновесие и макропеременные.
#
# Используется ЭНТРОПИЙНОЕ равновесие в product-form (как в KBC), а не усечённый полином:
#   feq_i = w_i · ρ · Π_d (2 − sqrt(1+3 u_d²)) · ((2 u_d + sqrt(1+3 u_d²))/(1 − u_d))^{c_{i,d}}
# Показатели c_{i,d} ∈ {−1,0,+1}, поэтому степень НЕ нужна — берём q или 1/q (быстрее и совместимо
# с CUDA-graph: нет pow и нет host→device для границ clip).
import backend as B
import lattice as L

cp = B.cp; xp = B.xp
CXa = L.CXa; CYa = L.CYa

def _feq9_body(rho, ux, uy):
    ux = cp.minimum(cp.maximum(ux, -0.95), 0.95)         # clip скоростей литералами (запекается в кернел)
    uy = cp.minimum(cp.maximum(uy, -0.95), 0.95)
    sx = cp.sqrt(1.0 + 3.0*ux*ux); sy = cp.sqrt(1.0 + 3.0*uy*uy)
    base = rho * (2.0 - sx) * (2.0 - sy)
    qx = (2.0*ux + sx) / (1.0 - ux); qy = (2.0*uy + sy) / (1.0 - uy)
    ix = 1.0/qx; iy = 1.0/qy
    return (base*(4.0/9.0),                               # ( 0, 0)
            base*(1.0/9.0)*qx,  base*(1.0/9.0)*qy,        # (+1,0) (0,+1)
            base*(1.0/9.0)*ix,  base*(1.0/9.0)*iy,        # (-1,0) (0,-1)
            base*(1.0/36.0)*qx*qy, base*(1.0/36.0)*ix*qy, # (+1,+1) (-1,+1)
            base*(1.0/36.0)*ix*iy, base*(1.0/36.0)*qx*iy) # (-1,-1) (+1,-1)

_feq9 = cp.fuse()(_feq9_body)        # одно слитное ядро (если cp.fuse доступен)
_FUSE_OK = [True]

def feq(rho, u):
    """Равновесные популяции (Q,Ny,Nx) по плотности rho и скорости u (индексируется [0],[1]).
    Сборка через B.pack (без cupy.stack) — иначе H2D ломает CUDA-graph capture."""
    if _FUSE_OK[0]:
        try:
            return B.pack(_feq9(rho, u[0], u[1]))
        except Exception as e:           # откат на не-fused (тоже без pow), один раз
            print(f"[equilibrium] cp.fuse отключён ({type(e).__name__}: {e}); обычный путь")
            _FUSE_OK[0] = False
    return B.pack(_feq9_body(rho, u[0], u[1]))

def macros(f):
    """Макропеременные: ρ = Σ_i f_i, u = (Σ_i c_i f_i)/ρ. Возвращает (rho, u=(2,Ny,Nx))."""
    r = f.sum(0)
    ux = (CXa[:, None, None] * f).sum(0) / r
    uy = (CYa[:, None, None] * f).sum(0) / r
    return r, B.pack((ux, uy))