# НЕЗАВИСИМЫЙ КРОСС-ЧЕК силы/момента: баланс импульса контрольного объёма, ограниченного
# замкнутым контуром (окружностью) вокруг тела на тонком уровне L1.
#
#   F_тела = ∮ [σ·n − ρ·u·(u·n)] ds − d/dt ∫_кольцо ρu dV,   σ = −p′·I + σ^v
#
#   * давление: p = ρ·c_s²; интегрируем p′ = (ρ − ⟨ρ⟩_контур)·c_s² — константа давления по замкнутому
#     контуру даёт ∮ p0·n ds ≡ 0, а вычитание убирает катастрофическое сокращение в float32;
#   * вязкая часть из неравновесных моментов (Чепмен–Энског): σ^v_αβ = −(1 − 1/(2τ))·Π^neq_αβ,
#     Π^neq_αβ = Σ_i c_iα c_iβ (f_i − f_i^eq).
#     ВАЖНО (KBC): след Π^neq релаксирует не с 1/τ, а с γβ (поле-зависимый стабилизатор), поэтому
#     берём строго ДЕВИАТОРНУЮ проекцию (след исключён) — сдвиг-моменты в KBC релаксируют с точным
#     2β = 1/τ, для них префактор корректен. Bulk-часть в силу трения и не входит.
#   * КОНВЕКТИВНЫЙ член −ρu(u·n): поток импульса через контур. Зануляется только при δ→0 (no-slip
#     на поверхности); при δ=2 скорость на контуре уже заметна, и без этого члена кросс-чек давал
#     систематический сдвиг +4% по ⟨Cd⟩ на всех сетках (факторное исследование, эксперимент №0).
#   * член d/dt ∫ρu dV по кольцу R..R1+δ НЕ считается: в среднем по времени он зануляется,
#     на мгновенных рядах даёт малый сдвиг фазы → кросс-чек ведётся по средним (⟨Cd⟩, rms Cl).
#   * контур — окружность радиуса R1+δ (δ = cfg.stress_delta тонких ячеек, по умолчанию 2):
#     билинейный стенсиль выборки не задевает первое жидкое кольцо у стенки, где популяции
#     реконструированы Bouzidi (q_max ≈ 0.98 < 1).
#
# Вызывается РЕДКО (каждые cfg.stress_every шагов) в Solver.run() ВНЕ step()/CUDA-графа.
import math
import backend as B
import lattice as L
import equilibrium as E

cp = B.cp; xp = B.xp; DTYPE = B.DTYPE
CS2 = L.CS2

def build_probe(cfg):
    """Прекомпьют контура (один раз, вне захвата): N=stress_ntheta точек на окружности R1+δ
    вокруг центра тела на L1; индексы/веса билинейной интерполяции, нормали, плечи, элемент дуги."""
    r = cfg.refine
    R1 = r * float(cfg.D) / 2.0                         # радиус цилиндра на L1 (тонкие ячейки)
    Rp = R1 + float(cfg.stress_delta)
    ccx = (cfg.cx - cfg.ax1) * r; ccy = (cfg.cy - cfg.ay1) * r
    N = int(cfg.stress_ntheta)
    th = (2.0 * math.pi / N) * cp.arange(N)
    nx = cp.cos(th).astype(DTYPE); ny = cp.sin(th).astype(DTYPE)
    px = ccx + Rp * nx; py = ccy + Rp * ny
    x0 = cp.floor(px).astype(cp.int64); y0 = cp.floor(py).astype(cp.int64)
    # контур обязан лежать внутри патча L1 (он центрирован вокруг тела — выполняется с запасом)
    Nfx = (cfg.bx1 - cfg.ax1) * r + 1; Nfy = (cfg.by1 - cfg.ay1) * r + 1
    assert int(x0.min()) >= 0 and int(x0.max()) + 1 <= Nfx - 1, "контур σ·n вышел за патч L1 по x"
    assert int(y0.min()) >= 0 and int(y0.max()) + 1 <= Nfy - 1, "контур σ·n вышел за патч L1 по y"
    return dict(x0=x0, y0=y0, x1=x0 + 1, y1=y0 + 1,
                tx=(px - x0).astype(DTYPE), ty=(py - y0).astype(DTYPE),
                nx=nx, ny=ny, rx=(Rp * nx), ry=(Rp * ny), ds=2.0 * math.pi * Rp / N)

def _sample(fld, pr):
    """Билинейная выборка 2D-поля в точках контура (аналог amr.pint, но по 1D-набору точек)."""
    return (fld[pr["y0"], pr["x0"]] * (1 - pr["tx"]) * (1 - pr["ty"])
            + fld[pr["y0"], pr["x1"]] * pr["tx"] * (1 - pr["ty"])
            + fld[pr["y1"], pr["x0"]] * (1 - pr["tx"]) * pr["ty"]
            + fld[pr["y1"], pr["x1"]] * pr["tx"] * pr["ty"])

def force_stress(F1, pr, tau1):
    """Сила/момент на тело из баланса импульса ∮[σ·n − ρu(u·n)]ds по контуру pr на состоянии F1
    (L1, после стриминга). Возвращает (Fx, Fy, Tz) — device-скаляры. tau1 — релаксация на L1."""
    rho, u = E.macros(F1)
    fneq = F1 - E.feq(rho, u)
    CX = L.CXa[:, None, None]; CY = L.CYa[:, None, None]
    Pxx = (CX * CX * fneq).sum(0)
    Pyy = (CY * CY * fneq).sum(0)
    Pxy = (CX * CY * fneq).sum(0)
    rho_s = _sample(rho, pr)
    ux_s = _sample(u[0], pr); uy_s = _sample(u[1], pr)
    Pxx_s = _sample(Pxx, pr); Pyy_s = _sample(Pyy, pr); Pxy_s = _sample(Pxy, pr)
    p = (rho_s - rho_s.mean()) * CS2                    # p′: давление без константы (см. шапку)
    pref = 1.0 - 1.0 / (2.0 * tau1)
    tr = 0.5 * (Pxx_s + Pyy_s)                          # девиаторная проекция (след исключён)
    sxx = -pref * (Pxx_s - tr); syy = -pref * (Pyy_s - tr); sxy = -pref * Pxy_s
    un = ux_s * pr["nx"] + uy_s * pr["ny"]              # нормальная скорость на контуре
    flux = rho_s * un                                   # ρ(u·n): конвективный вынос импульса
    tx = (-p + sxx) * pr["nx"] + sxy * pr["ny"] - flux * ux_s   # t = σ·n − ρu(u·n)
    ty = sxy * pr["nx"] + (-p + syy) * pr["ny"] - flux * uy_s
    Fx = tx.sum() * pr["ds"]; Fy = ty.sum() * pr["ds"]
    Tz = (pr["rx"] * ty - pr["ry"] * tx).sum() * pr["ds"]   # давление в торк не входит (n ∥ r)
    return Fx, Fy, Tz