# Граничные условия: (1) тело/стенки — SDF + интерполированный отскок Bouzidi; (2) вход/выход канала.
#
# Bouzidi (линейный): для линка из жидкого узла x_f в твёрдый по направлению i вводится
# доля пересечения q = |x_f→стенка| / |x_f→x_solid| ∈ (0,1):
#   q < 1/2 (есть «дальний» жидкий сосед x_f − c_i):  f_ī = 2q·f_i^post + (1−2q)·f_i^post(дальний)
#   q ≥ 1/2:                                           f_ī = (1/2q)·f_i^post + (1−1/2q)·f_ī^post
# Здесь f_i^post — после столкновения, f_ī — то, что вернётся в x_f после переноса.
import backend as B
import lattice as L

cp = B.cp; xp = B.xp; DTYPE = B.DTYPE
Q = L.Q; Cx = L.Cx; Cy = L.Cy; OPP = L.OPP

def build_bc(solid, phi, body):
    """Сборка структуры ГУ (один раз). Для каждого направления i (1..8):
      mask — жидкий узел, чей сосед по +c_i твёрдый;
      q    — доля пересечения по SDF (клип [0.02,0.98]);
      xff  — есть ли «дальний» жидкий сосед (для линейной формулы при q<1/2);
      cl   — подмаска линков именно ТЕЛА body (для расчёта силы)."""
    fluid = ~solid; bc = []
    for i in range(1, Q):
        nb_solid = cp.roll(solid, (-Cy[i], -Cx[i]), (0, 1)); mask = fluid & nb_solid
        cl = mask & cp.roll(body, (-Cy[i], -Cx[i]), (0, 1))
        phinb = cp.roll(phi, (-Cy[i], -Cx[i]), (0, 1))
        qq = phi / (phi - phinb)                                  # деление; nan/inf маскируются ниже
        q = xp.where(mask, cp.clip(qq, 0.02, 0.98), cp.asarray(0.5, DTYPE))
        xff = mask & (~cp.roll(solid, (Cy[i], Cx[i]), (0, 1)))
        bc.append((i, OPP[i], mask, q.astype(DTYPE), xff, cl))
    return bc

def apply_bc(f, fpost, bc):
    """Применить Bouzidi ко всем линкам (полностью через where, GPU-friendly).
    f — поле ПОСЛЕ стриминга (его и правим); fpost — поле ПОСЛЕ столкновения (до стриминга)."""
    for (i, ib, mask, q, xff, cl) in bc:
        fi = fpost[i]; fib = fpost[ib]
        fiback = cp.roll(fpost[i], (Cy[i], Cx[i]), (0, 1))        # популяция «дальнего» соседа
        near = mask & (q < 0.5) & xff
        bad  = mask & (q < 0.5) & (~xff)                         # нет дальнего соседа → простой отскок
        far  = mask & (q >= 0.5)
        new = f[ib]
        new = xp.where(near, 2*q*fi + (1 - 2*q)*fiback, new)
        new = xp.where(far,  (1/(2*q))*fi + (1 - 1/(2*q))*fib, new)
        new = xp.where(bad,  fi, new)
        f[ib] = new
    return f

def channel_bc(f, uin, rho_out=1.0, outlet_uy="zero"):
    """Корректные ГУ канала (Zou & He 1997), нумерация D2Q9: 1=E,2=N,3=W,4=S,5=NE,6=NW,7=SW,8=SE.
    ВХОД (запад, x=0): скоростной Zou-He — задаём u=(ux,uy) из uin, плотность ρ ПЛАВАЕТ; достраиваем
      приходящие извне популяции 1,5,8 из известных 0,2,3,4,6,7.
    ВЫХОД (восток, x=-1): давление Zou-He — задаём ρ=rho_out, скорость ux плавает; достраиваем 3,6,7.
      Поперечная скорость на выходе — по outlet_uy:
        "zero"        — классический Zou-He: uy=0 (жёстко). Вихри дорожки приходят на выход с
                        uy~±0.3U → принудительное зануление отражает их назад к телу (завышает
                        rms Cl и модулирует Cd — см. факторное исследование в ноутбуке);
        "extrapolate" — uy берётся из соседнего столбца x=-2 (нуль-градиент поперечной скорости);
                        ρ-якорь (масса) сохраняется, кинематическое отражение снимается.
    Зачем давление-выход вообще: грубый zero-gradient по f не выпускал поток → масса/противодавление
    копились, поток глох (⟨ρ⟩↑, Cd/rmsCl→0). Давление-выход якорит массу, скоростной вход её не
    пересоздаёт → дрейфа нет. Всё на срезах (graph-safe, без host→device).

    ПОРЯДОК ВЫЗОВА: строго ПОСЛЕ free_slip_walls и покрывая ВЕСЬ столбец (включая ряды 0 и Ny−1).
    Причина: stream() сделан через roll и потому периодичен по ОБЕИМ осям. Заворот по y снимают
    стенки, заворот по x — этот ГУ. Если оставить углы (ряды 0 и Ny−1) без Zou-He, в них популяции
    с cx=+1 на входе приходят прямо из столбца выхода (замерено: f[1,0,0] == post[1,0,-1], величина
    ~0.12) — вход и выход оказываются физически связаны в 4 узлах. Известные популяции, по которым
    здесь считаются ρ (вход) и ux (выход), x-заворота не содержат ни в одном ряду, а те, что
    заворот содержат, — это ровно неизвестные, которые ГУ и перезаписывает; поэтому такой порядок
    закрывает углы полностью."""
    # --- вход (запад): скоростной Zou-He ---
    ux = uin[0, :, 0]; uy = uin[1, :, 0]
    f0 = f[0, :, 0]; f2 = f[2, :, 0]; f3 = f[3, :, 0]
    f4 = f[4, :, 0]; f6 = f[6, :, 0]; f7 = f[7, :, 0]
    rho = (f0 + f2 + f4 + 2.0*(f3 + f6 + f7)) / (1.0 - ux)
    f[1, :, 0] = f3 + (2.0/3.0)*rho*ux
    f[5, :, 0] = f7 - 0.5*(f2 - f4) + (1.0/6.0)*rho*ux + 0.5*rho*uy
    f[8, :, 0] = f6 + 0.5*(f2 - f4) + (1.0/6.0)*rho*ux - 0.5*rho*uy
    # --- выход (восток): давление Zou-He ---
    g0 = f[0, :, -1]; g1 = f[1, :, -1]; g2 = f[2, :, -1]
    g4 = f[4, :, -1]; g5 = f[5, :, -1]; g8 = f[8, :, -1]
    uxo = (g0 + g2 + g4 + 2.0*(g1 + g5 + g8)) / rho_out - 1.0
    if outlet_uy == "extrapolate":
        # нуль-градиент поперечной скорости: uy с предвыходного столбца (он полностью известен)
        col = f[:, :, -2]
        uyo = ((col[2] + col[5] + col[6]) - (col[4] + col[7] + col[8])) / col.sum(0)
    else:
        uyo = 0.0
    f[3, :, -1] = g1 - (2.0/3.0)*rho_out*uxo
    f[6, :, -1] = g8 - 0.5*(g2 - g4) - (1.0/6.0)*rho_out*uxo + 0.5*rho_out*uyo
    f[7, :, -1] = g5 + 0.5*(g2 - g4) - (1.0/6.0)*rho_out*uxo - 0.5*rho_out*uyo
    return f

def free_slip_walls(f):
    """FREE-SLIP (скользящие) стенки канала: верх/низ — зеркальное (specular) отражение.
    Сохраняет касательный (x) импульс (нет трения, нет пограничного слоя) и нормальный заворачивает;
    масса сохраняется (копирование популяций) → дрейфа не вносит. Применять ПОСЛЕ стриминга.
    Нумерация: 2=(0,+1) 4=(0,-1) 5=(+1,+1) 6=(-1,+1) 7=(-1,-1) 8=(+1,-1).
    Нижняя стенка (row 0): достраиваем вверх-идущие из вниз-идущих (зеркало по y, x сохранён).
    Верхняя стенка (row -1): достраиваем вниз-идущие из вверх-идущих.
    ВНИМАНИЕ: это только искусственные стенки канала; цилиндр остаётся no-slip (Bouzidi).
    ПОРЯДОК: вызывать ПЕРЕД channel_bc (см. его docstring — так закрываются углы домена)."""
    f[2, 0, :] = f[4, 0, :]; f[5, 0, :] = f[8, 0, :]; f[6, 0, :] = f[7, 0, :]
    f[4, -1, :] = f[2, -1, :]; f[7, -1, :] = f[6, -1, :]; f[8, -1, :] = f[5, -1, :]
    return f