# Оператор СТОЛКНОВЕНИЯ — KBC-N1 (энтропийный), единственный оператор модели.
#
# KBC-N1 (Karlin/Bösch/Chikatamarla 2014):
#   Δ   = f − feq                           (неравновесная часть)
#   Δs  = Ps·Δ                              (проекция на сдвиг-моменты {cx²−cy², cxcy})
#   Δh  = Δ − Δs                            (остальная часть)
#   γ   = 1/β − (2 − 1/β)·⟨Δs,Δh⟩ / ⟨Δh,Δh⟩,   ⟨a,b⟩ = Σ_i a_i b_i / feq_i  (H-теорема)
#   f_post = f − β(2Δs + γΔh),   β = 1/(2τ)
# γ — энтропийный стабилизатор: подстраивает диссипацию неhydro-мод под энтропийное ограничение
# (поэтому KBC устойчив при τ→1/2). Сдвиг-моменты (вязкость) релаксируют с правильным 2β = 1/τ,
# то есть физическая вязкость задаётся только τ — стабилизатор её не меняет.
#
# Порядок шагов дословно совпадает с алгоритмом статьи (Bösch/Chikatamarla/Karlin, PRE 92,
# 043309 (2015), разд. III, листинг после ур.39): ρ,u → feq → Δs → Δh = f − feq − Δs → γ → релаксация.
# Вырожденный случай ⟨Δh|Δh⟩→0 отсекается ОТНОСИТЕЛЬНЫМ порогом B.GREL (см. backend.py):
# абсолютный порог здесь недопустим — den квадратична по неравновесию и физически мала.
import backend as B
import lattice as L

cp = B.cp; xp = B.xp
Q = L.Q; Ps = L.Ps; GREL = B.GREL; ONE = B.ONE; TWO = B.TWO

def _project_shift(dfn):
    """Δs = Ps·Δ без cuBLAS (gemm при CUDA-graph capture тащит alpha/beta host→device).
    Реализовано как Σ_k Ps[i,k]·Δ[k] через broadcast+редукцию — только обычные ядра."""
    dfn2d = dfn.reshape(Q, -1)                                    # (Q, N)
    return (Ps[:, :, None] * dfn2d[None]).sum(1).reshape(dfn.shape)

def kbc_collide(f, fe, beta):
    dfn = f - fe
    ds = _project_shift(dfn)
    dh = dfn - ds; inv = 1.0 / fe
    num = (ds*dh*inv).sum(0); den = (dh*dh*inv).sum(0)
    nrm = (dfn*dfn*inv).sum(0)                                   # ‖Δ‖² — масштаб неравновесия узла
    ok = den > GREL*nrm                                          # относительный порог, не абсолютный
    den_safe = xp.where(ok, den, ONE)                            # избегаем 0/0 (обе ветки where считаются)
    binv = 1.0 / beta
    g = xp.where(ok, binv - (2.0 - binv)*num/den_safe, TWO)
    return f - beta*(2.0*ds + g[None]*dh)

# единственный оператор столкновения модели
collide = kbc_collide