# Самопроверки математики НА КОМПОНЕНТАХ ФИНАЛЬНОЙ МОДЕЛИ (solver_2x_sdf): решётка,
# энтропийное равновесие, KBC-проектор, предел γ=2 ≡ BGK. Результаты → out/checks.txt
# (ноутбук цитирует их в markdown). Запуск: python checks_gpu.py [AMR_FP64=1 — жёстче пороги].
import numpy as np
import _common as cm

cp = cm.cp; xp = cm.xp; L = cm.L; E = cm.E; Coll = cm.Coll
Q = cm.Q; CS2 = cm.CS2; TOL = cm.TOL_EXACT

lines = [f"backend={cm.B.BACKEND}; порог точных тождеств TOL={TOL:g}"]
def check(name, value, tol):
    ok = value < tol
    lines.append(f"{'PASS' if ok else 'FAIL'}  {name}: {value:.3e} (tol {tol:g})")
    print(lines[-1])
    assert ok, name
    return ok

def check_gt(name, value, floor):
    """Для величин, которые обязаны быть НЕ малы (иначе модель выродилась)."""
    ok = value > floor
    lines.append(f"{'PASS' if ok else 'FAIL'}  {name}: {value:.3e} (должно быть > {floor:g})")
    print(lines[-1])
    assert ok, name
    return ok

# --- 1. Решётка D2Q9: тождества весов и изотропии ---------------------------------------------
# Допуск по типу: Wa хранится в DTYPE, и в float32 уже само Σw отличается от 1 на ~7e-9,
# так что жёсткий 1e-12 здесь означал бы падение проверки в режиме по умолчанию.
TOL_W = 1e-12 if cm.FP64 else 1e-6
C = np.array(list(zip(L.Cx, L.Cy)), float)
W = cm.to_cpu(L.Wa).astype(float)
check("Σw − 1", abs(W.sum() - 1.0), TOL_W)
check("|Σ w·c|", np.abs((W[:, None]*C).sum(0)).max(), TOL_W)
check("|Σ w·c⊗c − cs²·I|",
      np.abs((W[:, None, None]*C[:, :, None]*C[:, None, :]).sum(0) - CS2*np.eye(2)).max(), TOL_W)
assert np.all(C[list(L.OPP)] == -C), "OPP: c_opp != -c"
lines.append("PASS  OPP: c_opp = −c (точно)")

# --- 2. Энтропийное равновесие (product-form) точно воспроизводит ρ и ρu -----------------------
rng = np.random.default_rng(0)
rho_t = xp.asarray(1.0 + 0.1*rng.standard_normal((16, 16)), cm.DTYPE)
u_t = xp.asarray(0.06*rng.standard_normal((2, 16, 16)), cm.DTYPE)
fe = E.feq(rho_t, u_t)
rho_m, u_m = E.macros(fe)
check("feq: |Σf − ρ|", float(xp.abs(rho_m - rho_t).max()), TOL)
check("feq: |u(f) − u|", float(xp.abs(u_m - u_t).max()), TOL)
Pxx = (L.CXa[:, None, None]**2 * fe).sum(0)
check("feq: |Σc_x²f − ρ(cs²+u_x²)| (мал при малом Ma)",
      float(xp.abs(Pxx - rho_t*(CS2 + u_t[0]**2)).max()), 5e-3)

# --- 3. KBC-проектор Ps: идемпотентность и состав сдвиг-моментов --------------------------------
Ps = cm.to_cpu(L.Ps).astype(float)
check("Ps идемпотентен: |Ps·Ps − Ps|", np.abs(Ps @ Ps - Ps).max(), 1e-5 if not cm.FP64 else 1e-12)
# моментная матрица (как в lattice._build_projector): Δs несёт ТОЛЬКО {cx²−cy², cx·cy}
cx, cy = C[:, 0], C[:, 1]
M = np.zeros((Q, Q))
M[0] = 1.0; M[1] = cx; M[2] = cy; M[3] = 3*(cx**2 + cy**2) - 2
M[4] = cx**2 - cy**2; M[5] = cx*cy
M[6] = cx**2*cy; M[7] = cx*cy**2; M[8] = cx**2*cy**2
dfn = rng.standard_normal((Q, 5, 7))
ds = np.einsum("ij,jab->iab", Ps, dfn)
mds = np.einsum("ai,ijk->ajk", M, ds)
other = [0, 1, 2, 3, 6, 7, 8]
check("Ps: |моменты Δs вне {N, Πxy}|", np.abs(mds[other]).max(), 1e-4 if not cm.FP64 else 1e-10)
lines.append("PASS  состав KBC-N1: s = {Πxx−Πyy, Πxy}, h = остальное (k отщеплён точно)")

# --- 4. Предел γ=2: KBC-форма с γ=2 совпадает с BGK (через РЕАЛЬНЫЙ проектор модели) ------------
f0 = fe + xp.asarray(0.01*rng.standard_normal((Q, 16, 16)), cm.DTYPE)
r0, u0 = E.macros(f0)
fe0 = E.feq(r0, u0)
beta = 0.7
dfn0 = f0 - fe0
ds0 = Coll._project_shift(dfn0)
fk2 = f0 - beta*(2.0*ds0 + 2.0*(dfn0 - ds0))     # KBC с γ=2
fbgk = cm.bgk_collide(f0, fe0, beta)
check("KBC(γ=2) ≡ BGK", float(xp.abs(fk2 - fbgk).max()), TOL)

# --- 5. kbc_collide сохраняет ρ и импульс -------------------------------------------------------
fpost = Coll.kbc_collide(f0, fe0, beta)
r1, u1 = E.macros(fpost)
check("KBC: |Δρ| за столкновение", float(xp.abs(r1 - r0).max()), TOL)
check("KBC: |Δ(ρu)| за столкновение", float(xp.abs(u1*r1[None] - u0*r0[None]).max()), TOL)

# --- 6. Стабилизатор γ ДЕЙСТВИТЕЛЬНО считается на физическом поле ------------------------------
# Регрессия на дефект от 2026-08-14: абсолютный порог знаменателя (1e-6 в fp32) отправлял в откат
# γ←2 77–99% узлов, т.е. модель молча вырождалась в LBGK. Проверяем на РАЗВИТОМ поле (сдвиговый
# слой после прогрева), что откат не срабатывает и γ не залипает на 2.
Strm = cm.Strm
Nsl = 64; u0sl = 0.04; beta_sl = cm.nu_to_beta(u0sl*Nsl/5000.0)
yy = xp.arange(Nsl, dtype=cm.DTYPE)[:, None] * xp.ones((1, Nsl), cm.DTYPE)
xx = xp.ones((Nsl, 1), cm.DTYPE) * xp.arange(Nsl, dtype=cm.DTYPE)[None, :]
uxsl = xp.where(yy <= Nsl/2, u0sl*xp.tanh(40.0*(yy/Nsl - 0.25)), u0sl*xp.tanh(40.0*(0.75 - yy/Nsl)))
uysl = 0.05*u0sl*xp.sin(2.0*np.pi*(xx/Nsl + 0.25))
fsl = E.feq(xp.ones((Nsl, Nsl), cm.DTYPE), cm.B.pack((uxsl, uysl)))
for _ in range(400):
    rsl, usl = E.macros(fsl)
    fsl = Strm.stream(Coll.kbc_collide(fsl, E.feq(rsl, usl), beta_sl))
rsl, usl = E.macros(fsl); fesl = E.feq(rsl, usl)
dfs = fsl - fesl; dss = Coll._project_shift(dfs); dhs = dfs - dss; invs = 1.0/fesl
den_s = (dhs*dhs*invs).sum(0); nrm_s = (dfs*dfs*invs).sum(0)
frac_fb = float((den_s <= cm.B.GREL*nrm_s).mean())
check("γ: доля узлов в откате γ←2 на развитом поле", frac_fb, 1e-3)
gam = cm.kbc_gamma(fsl, fesl, beta_sl)
# при полном вырождении в LBGK эта величина строго 0, при частичном (90% узлов) — ~0.01,
# наблюдаемое значение ~0.13, поэтому порог 0.05 ловит регресс с запасом и не флапает
check_gt("γ: ⟨|γ−2|⟩ — отличие от LBGK", float(xp.abs(gam - 2.0).mean()), 0.05)
lines.append(f"INFO  γ на развитом поле: ⟨γ⟩={float(gam.mean()):.4f} "
             f"[{float(gam.min()):.3f}, {float(gam.max()):.3f}], min den/‖Δ‖²="
             f"{float((den_s/xp.maximum(nrm_s, 1e-30)).min()):.2e} (порог {cm.B.GREL:g})")
print(lines[-1])

# --- 7. Замкнутая формула γ ур.(17)/(39) против корня точного ур.(15) --------------------------
# Берём один узел развитого поля и решаем Σ Δh·ln(1 + ((1−βγ)Δh − (2β−1)Δs)/feq) = 0 бисекцией.
iy, ix = Nsl//3, Nsl//3
fe_n = cm.to_cpu(fesl[:, iy, ix]).astype(float)
ds_n = cm.to_cpu(dss[:, iy, ix]).astype(float); dh_n = cm.to_cpu(dhs[:, iy, ix]).astype(float)
g_formula = float(cm.to_cpu(gam[iy, ix]))
def _crit(g):
    return float((dh_n*np.log(1.0 + ((1.0 - beta_sl*g)*dh_n - (2.0*beta_sl - 1.0)*ds_n)/fe_n)).sum())
lo, hi = g_formula - 2.0, g_formula + 2.0
while _crit(lo)*_crit(hi) > 0 and hi - lo < 40:
    lo -= 1.0; hi += 1.0
for _ in range(200):
    mid = 0.5*(lo + hi)
    if _crit(lo)*_crit(mid) <= 0: hi = mid
    else: lo = mid
check("γ: |формула − точный корень ур.(15)|", abs(g_formula - 0.5*(lo + hi)), 5e-2)

# --- 8. Гидродинамический предел: ν по затуханию сдвиговой волны --------------------------------
# Единственная проверка, которая ловит подмену оператора: LBGK и KBC дают одну ν, но неверная
# релаксация сдвиг-моментов сразу видна здесь.
Nw = 48; kw = 2.0*np.pi/Nw; NSW = 1500
sinw = xp.sin(kw*xp.arange(Nw, dtype=cm.DTYPE))
for tau_w in (0.6, 1.0):
    beta_w = 1.0/(2.0*tau_w); nu_th = CS2*(tau_w - 0.5)
    # амплитуда 1e-2 (Ma≈0.017, режим ещё линейный): в float32 при 1e-4 сигнал тонет в округлении
    uxw = 0.01*sinw[:, None]*xp.ones((1, 4), cm.DTYPE)
    fw = E.feq(xp.ones((Nw, 4), cm.DTYPE), cm.B.pack((uxw, xp.zeros_like(uxw))))
    amps = []
    for _ in range(NSW + 1):
        rw, uw = E.macros(fw)
        amps.append(abs(float((uw[0][:, 0]*sinw).sum())))      # проекция на моду, без abs поэлементно
        fw = Strm.stream(Coll.kbc_collide(fw, E.feq(rw, uw), beta_w))
    a = np.array(amps); i0 = NSW//3
    rate = -np.polyfit(np.arange(i0, NSW + 1), np.log(a[i0:]), 1)[0]
    # допуск 1e-2 покрывает решёточную дисперсию O(k²) (k=0.131 → ~3e-3)
    check(f"ν(τ={tau_w}): отн. ошибка затухания сдвиговой волны", abs(rate/(kw*kw)/nu_th - 1.0), 1e-2)

# --- 9. ГУ канала: вход и выход не связаны заворотом roll (включая углы) ------------------------
# Регрессия на дефект от 2026-08-14: Zou-He стоял на срезе 1:-1, и в 4 углах stream() через roll
# сносил популяции прямо из столбца выхода в столбец входа.
import boundary as Bnd
Ny_b, Nx_b = 24, 32
fb = E.feq(xp.ones((Ny_b, Nx_b), cm.DTYPE), 0.02*xp.ones((2, Ny_b, Nx_b), cm.DTYPE))
fb = fb*(1.0 + xp.asarray(0.05*rng.standard_normal(fb.shape), cm.DTYPE))
post_b = fb.copy()
uin_b = xp.zeros((2, Ny_b, 1), cm.DTYPE); uin_b[0] = 0.05
fs_b = Bnd.channel_bc(Bnd.free_slip_walls(Strm.stream(post_b)), uin_b, outlet_uy="zero")
r_in = fs_b[:, :, 0].sum(0)
ux_in = (L.CXa[:, None]*fs_b[:, :, 0]).sum(0)/r_in
check("ГУ: |u_x(вход) − заданной| по ВСЕМ рядам, вкл. углы", float(xp.abs(ux_in - 0.05).max()), TOL)
check("ГУ: |ρ(выход) − 1| по ВСЕМ рядам, вкл. углы",
      float(xp.abs(fs_b[:, :, -1].sum(0) - 1.0).max()), TOL)
leak = min(float(xp.abs(fs_b[i, row, 0] - post_b[i, row, -1]).max())
           for i in (1, 5, 8) for row in (0, -1))
check_gt("ГУ: расхождение углового f с популяцией выхода (заворот подавлен)", leak, 1e-6)

lines.append("=== Все самопроверки пройдены ===")
print(lines[-1])
cm.write_out("checks.txt", lines)