# Общая инфраструктура GPU-демо (ноутбук kbc_lbm.ipynb код не исполняет — только встраивает
# артефакты, которые генерят эти скрипты на GPU-сервере).
#
# Физика НЕ дублируется: решётка/равновесие/столкновение/перенос импортируются из
# ../solver_2x_sdf (те же модули, что и в финальной модели) — демо тем самым доказывают,
# что КОМПОНЕНТЫ финальной модели воспроизводят классические бенчмарки.
# Здесь только то, чего в модели нет намеренно: полиномиальное равновесие и BGK
# (для сравнений «BGK vs KBC»), H-функция, γ-поле для визуализации, гифки.
import os
import sys
import io as _io

HERE = os.path.dirname(os.path.abspath(__file__))
THEORY = os.path.dirname(HERE)
sys.path.insert(0, os.path.join(THEORY, "solver_2x_sdf"))

import backend as B          # GPU-only: жёсткая ошибка без CuPy (CPU-фолбэка нет намеренно)
import lattice as L
import equilibrium as E
import collision as Coll
import streaming as Strm
import visualization as Viz  # make_progress (ASCII-бар)

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

import numpy as np

# пути артефактов: фигуры/гифки — общие каталоги theory, числа — demos_gpu/out
FIGDIR = os.path.join(THEORY, "figures"); os.makedirs(FIGDIR, exist_ok=True)
ANIMDIR = os.path.join(THEORY, "anim"); os.makedirs(ANIMDIR, exist_ok=True)
OUTDIR = os.path.join(HERE, "out"); os.makedirs(OUTDIR, exist_ok=True)

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

try:
    sys.stdout.reconfigure(encoding="utf-8")      # кириллица/греческие в print при cp1251-консоли
except Exception:
    pass

plt.rcParams.update({
    "figure.dpi": 110, "savefig.dpi": 120, "font.size": 11,
    "axes.titlesize": 13, "axes.titleweight": "bold",
    "axes.grid": True, "grid.alpha": 0.25, "axes.axisbelow": True,
    "figure.facecolor": "white", "axes.facecolor": "#fbfbfd",
})

# пороги самопроверок под точность бэкенда (демо обычно гоняются в float32)
FP64 = (DTYPE == cp.float64)
TOL_EXACT = 1e-12 if FP64 else 2e-5       # «точные» тождества (моменты равновесия и т.п.)

def finish(fig, name):
    """Сохранить фигуру в figures/ под тем же именем, что использовал ноутбук."""
    fig.savefig(os.path.join(FIGDIR, name), bbox_inches="tight")
    plt.close(fig)
    print("PNG:", name)

def write_out(name, lines):
    """Дописать результаты в demos_gpu/out/<name> (ноутбук цитирует эти числа в markdown)."""
    with open(os.path.join(OUTDIR, name), "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    print("OUT:", name)

# ----- дополнения к модели (нужны только демо) ------------------------------------------------

def feq_poly(rho, u):
    """Полиномиальное равновесие O(u²) — стандарт BGK (в модели НЕ используется; здесь —
    для демонстрации «BGK+poly разрушается там, где KBC+entropic держит»)."""
    CXa = L.CXa[:, None, None]; CYa = L.CYa[:, None, None]; Wa = L.Wa[:, None, None]
    cu = CXa*u[0] + CYa*u[1]
    usq = u[0]**2 + u[1]**2
    return Wa * rho * (1.0 + cu/CS2 + cu**2/(2.0*CS2*CS2) - usq/(2.0*CS2))

def bgk_collide(f, fe, beta):
    """Классический BGK: f − ω(f−feq), ω=2β. Эквивалент KBC при γ=2."""
    return f - 2.0*beta*(f - fe)

def kbc_gamma(f, fe, beta):
    """Поле энтропийного стабилизатора γ (для визуализации; сам оператор — collision.kbc_collide,
    формула та же: γ = 1/β − (2−1/β)·⟨Δs|Δh⟩/⟨Δh|Δh⟩)."""
    dfn = f - fe
    ds = Coll._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)                     # тот же относительный порог, что в collision.py
    ok = den > B.GREL*nrm
    return xp.where(ok, 1.0/beta - (2.0 - 1.0/beta)*num/xp.where(ok, den, B.ONE), B.TWO)

def nu_to_beta(nu):
    return 1.0 / (2.0 * (nu / CS2 + 0.5))

def H_function(f):
    """Дискретная H-функция H = Σ f_i ln(f_i/w_i) (поле)."""
    fp = xp.maximum(f, 1e-30)
    return (fp * xp.log(fp / L.Wa[:, None, None])).sum(0)

def smoothstep(x):
    x = min(max(x, 0.0), 1.0)
    return x*x*(3.0 - 2.0*x)

# ----- host-постобработка (numpy: кадры уже перенесены на хост) --------------------------------

def vorticity(u):
    """ω = ∂x u_y − ∂y u_x центральными разностями (host, по кадру)."""
    duy_dx = (np.roll(u[1], -1, 1) - np.roll(u[1], 1, 1)) * 0.5
    dux_dy = (np.roll(u[0], -1, 0) - np.roll(u[0], 1, 0)) * 0.5
    return duy_dx - dux_dy

def render_gif(n_frames, draw_fn, name, figsize, fps=18, dpi=80):
    """Отрисовать n_frames кадров (draw_fn(fig, i)) и сохранить зацикленный GIF в anim/."""
    from PIL import Image
    path = os.path.join(ANIMDIR, name)
    fig = plt.figure(figsize=figsize, dpi=dpi)
    pil = []
    for i in range(n_frames):
        fig.clf()
        draw_fn(fig, i)
        buf = _io.BytesIO()
        fig.savefig(buf, format="png")
        buf.seek(0)
        pil.append(Image.open(buf).convert("RGB").convert("P", palette=Image.ADAPTIVE))
    plt.close(fig)
    pil[0].save(path, save_all=True, append_images=pil[1:],
                duration=int(1000 / fps), loop=0, optimize=True)
    print("GIF:", name, f"({n_frames} кадров)")
    return path