# Статические иллюстрации (НЕ симуляции): стенсили D2Q9/D3Q27, блок-схема шага KBC,
# схема блочного измельчения. Чистый numpy+matplotlib — НЕ импортирует CuPy и потому
# исполним и локально (единственный скрипт demos_gpu без GPU).
# Схема измельчения использует фон |ω| из out/obstacle_circle_field.npz (его пишет
# obstacle_gpu.py на сервере); без него рисуется нейтральный фон.
# Артефакты: figures/01_stencil_d2q9.png, 06_algorithm_flow.png, 11c_refinement_schematic.png,
# 11_d3q27.png. Запуск: python figures_static.py
import os
import sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.lines import Line2D

try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass

HERE = os.path.dirname(os.path.abspath(__file__))
THEORY = os.path.dirname(HERE)
FIGDIR = os.path.join(THEORY, "figures"); os.makedirs(FIGDIR, exist_ok=True)
OUTDIR = os.path.join(HERE, "out")

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",
})

def finish(fig, name):
    fig.savefig(os.path.join(FIGDIR, name), bbox_inches="tight")
    plt.close(fig)
    print("PNG:", name)

# D2Q9 (та же нумерация, что в solver_2x_sdf/lattice.py)
CX = np.array([0, 1, 0, -1, 0, 1, -1, -1, 1], float)
CY = np.array([0, 0, 1, 0, -1, 1, 1, -1, -1], float)
W = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36])
CS2 = 1.0/3.0

# --- 1. Стенсиль D2Q9 ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(5.2, 5.2))
tier_color = {4/9: "#d1495b", 1/9: "#2e7d32", 1/36: "#1f6feb"}
for i in range(9):
    col = tier_color[W[i]]
    if i == 0:
        ax.plot(0, 0, "o", ms=18, color=col, zorder=3)
        ax.annotate(f"$w_0={W[i]:.3f}$", (0, 0), textcoords="offset points",
                    xytext=(8, 8), fontsize=9)
    else:
        ax.annotate("", xy=(CX[i], CY[i]), xytext=(0, 0),
                    arrowprops=dict(arrowstyle="-|>", lw=1.5 + 40*W[i], color=col))
        ax.annotate(f"$c_{{{i}}}$", (CX[i], CY[i]), textcoords="offset points",
                    xytext=(6*np.sign(CX[i] + 0.1), 6*np.sign(CY[i] + 0.1)), fontsize=9)
ax.set_xlim(-1.6, 1.6); ax.set_ylim(-1.6, 1.6); ax.set_aspect("equal")
ax.set_title("Решётка D2Q9: дискретные скорости $\\mathbf{c}_i$ и веса $w_i$")
ax.set_xlabel("$c_x$ [lu/ts]"); ax.set_ylabel("$c_y$ [lu/ts]")
ax.legend(handles=[Line2D([0], [0], color=c, lw=3, label=f"$w={k:.4f}$")
                   for k, c in tier_color.items()], loc="upper right", fontsize=8)
finish(fig, "01_stencil_d2q9.png")

# --- 2. Блок-схема шага KBC ---------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(4.3, 6.2))
ax.set_axis_off()
steps_txt = [
    ("Макро-моменты\n$\\rho=\\sum f_i,\\ \\rho\\mathbf{u}=\\sum \\mathbf{c}_i f_i$", "#e8f0fe"),
    ("Равновесие $f_i^{eq}$\n(энтропийное)", "#e6f4ea"),
    ("Неравновесие\n$\\Delta s=P_s(f-f^{eq}),\\ \\Delta h=(f-f^{eq})-\\Delta s$", "#fff4e5"),
    ("Стабилизатор $\\gamma$\n$=1/\\beta-(2-1/\\beta)\\frac{\\langle\\Delta s|\\Delta h\\rangle}{\\langle\\Delta h|\\Delta h\\rangle}$", "#fde7ef"),
    ("Столкновение\n$f\\leftarrow f-\\beta(2\\Delta s+\\gamma\\Delta h)$", "#f3e8fd"),
    ("Перенос\n$f_i(\\mathbf{x}+\\mathbf{c}_i)\\leftarrow f_i(\\mathbf{x})$", "#e8f0fe"),
    ("Граничные условия", "#eceff1"),
]
y = 0.95
for txt, col in steps_txt:
    ax.add_patch(plt.Rectangle((0.05, y - 0.1), 0.9, 0.085, transform=ax.transAxes,
                               facecolor=col, edgecolor="#555", lw=1.2))
    ax.text(0.5, y - 0.057, txt, transform=ax.transAxes, ha="center", va="center", fontsize=8.5)
    if y > 0.2:
        ax.annotate("", xy=(0.5, y - 0.118), xytext=(0.5, y - 0.1), xycoords="axes fraction",
                    arrowprops=dict(arrowstyle="-|>", color="#555"))
    y -= 0.13
ax.set_title("Шаг KBC")
finish(fig, "06_algorithm_flow.png")

# --- 3. Схема блочного измельчения (фон — |ω| цилиндра, если есть артефакт) ----------------------
def _vorticity(u):
    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

field_path = os.path.join(OUTDIR, "obstacle_circle_field.npz")
schematic_path = os.path.join(FIGDIR, "11c_refinement_schematic.png")
if not os.path.exists(field_path) and os.path.exists(schematic_path):
    # без серверного фона не затираем уже существующую схему с |ω|-подложкой
    print("(нет out/obstacle_circle_field.npz — оставляю существующий 11c_refinement_schematic.png; "
          "запустите obstacle_gpu.py на сервере и повторите)")
else:
    fig, ax = plt.subplots(figsize=(11, 5.6))
    if os.path.exists(field_path):
        d = np.load(field_path)
        u = d["u"]; solid = d["solid"]
        Nx, Ny = int(d["Nx"]), int(d["Ny"]); cx, cy, D = int(d["cx"]), int(d["cy"]), int(d["D"])
        vmag = np.where(solid, np.nan, np.abs(_vorticity(u)))
        ax.imshow(vmag, origin="lower", cmap="inferno", vmax=np.nanpercentile(vmag, 99))
    else:
        Nx, Ny, cx, cy, D = 150, 72, 37, 36, 14
        ax.add_patch(plt.Circle((cx, cy), D/2, color="#37474f"))
        ax.set_facecolor("#10101a")

    def _grid_block(x0, y0, w, h, step, color, label):
        ax.add_patch(mpatches.Rectangle((x0, y0), w, h, fill=False, edgecolor=color, lw=2.2))
        for xx in np.arange(x0, x0 + w + 0.1, step):
            ax.plot([xx, xx], [y0, y0 + h], color=color, lw=0.4, alpha=0.45)
        for yy in np.arange(y0, y0 + h + 0.1, step):
            ax.plot([x0, x0 + w], [yy, yy], color=color, lw=0.4, alpha=0.45)
        ax.text(x0 + 1, y0 + h - 4, label, color=color, fontsize=8, weight="bold")

    _grid_block(0, 0, Nx - 1, Ny - 1, 16, "#4fc3f7", "уровень 0: Δx (крупный)")
    _grid_block(cx - 1.5*D, cy - 1.9*D, 6.5*D, 3.8*D, 8, "#aed581", "уровень 1: Δx/2 (тело+след)")
    _grid_block(cx - 1.1*D, cy - 1.3*D, 2.6*D, 2.6*D, 4, "#ff8a65", "уровень 2: Δx/4 (у тела)")
    ax.set_xlim(0, Nx - 1); ax.set_ylim(0, Ny - 1)
    ax.set_title("Схема блочного измельчения вокруг тела (поверх $|\\omega|$): "
                 "сетка мельче там, где градиенты резче")
    ax.set_xlabel("x [lu]"); ax.set_ylabel("y [lu]")
    finish(fig, "11c_refinement_schematic.png")

# --- 4. Стенсиль D3Q27 --------------------------------------------------------------------------
C3 = np.array([[x, y, z] for x in (-1, 0, 1) for y in (-1, 0, 1) for z in (-1, 0, 1)], dtype=float)
nz = (C3 != 0).sum(1)
W3 = np.select([nz == 0, nz == 1, nz == 2, nz == 3], [8/27, 2/27, 1/54, 1/216])
assert np.isclose(W3.sum(), 1.0)
assert np.allclose((W3[:, None, None]*C3[:, :, None]*C3[:, None, :]).sum(0), CS2*np.eye(3))
print("D3Q27: Σw=1 и Σw c⊗c = cs²·I — OK; направлений:", len(C3))

fig = plt.figure(figsize=(6.4, 5.6))
ax = fig.add_subplot(111, projection="3d")
tier3 = {8/27: ("#d1495b", "покой"), 2/27: ("#2e7d32", "оси (×6)"),
         1/54: ("#1f6feb", "рёбра (×12)"), 1/216: ("#9c27b0", "углы (×8)")}
for wv, (col, lab) in tier3.items():
    m = np.isclose(W3, wv)
    ax.scatter(C3[m, 0], C3[m, 1], C3[m, 2], s=70 + 5000*wv, color=col,
               label=f"{lab}, w={wv:.4f}", depthshade=True, edgecolors="k", linewidths=0.4)
ax.set_xlabel("$c_x$"); ax.set_ylabel("$c_y$"); ax.set_zlabel("$c_z$")
ax.set_title("Решётка D3Q27 (цель реализации)")
ax.legend(loc="upper left", fontsize=7, bbox_to_anchor=(0.0, 1.0))
finish(fig, "11_d3q27.png")

print("DONE (static figures)")