# Демо-4: обтекание тел разной формы (цилиндр/квадрат/треугольник/профиль) при Re≈150 —
# вихревые дорожки Кармана на грубой сетке (KBC держит устойчиво). Ступенчатый bounce-back
# (узловой) — именно его «лесенку» позже лечит SDF+Bouzidi в финальной модели.
# Артефакты: figures/10_obstacle_shapes.png, figures/10b_obstacle_probe.png,
# anim/obstacle_shapes.gif, out/obstacle.txt, out/obstacle_circle_field.npz (фон для схемы AMR).
# Запуск: python obstacle_gpu.py
import os
import numpy as np
import _common as cm

cp = cm.cp; xp = cm.xp; E = cm.E; Coll = cm.Coll; Strm = cm.Strm
plt = cm.plt; vorticity = cm.vorticity; OPP = cm.OPP; Q = cm.Q

def solid_mask(Nx, Ny, shape, cx, cy, D):
    """Маска тела + стенки канала (host numpy → device bool)."""
    Y, X = np.mgrid[0:Ny, 0:Nx]
    if shape == "circle":
        m = (X - cx)**2 + (Y - cy)**2 <= (D/2)**2
    elif shape == "square":
        m = (np.abs(X - cx) <= D/2) & (np.abs(Y - cy) <= D/2)
    elif shape == "triangle":          # вершиной против потока
        tt = (X - (cx - D/2))/D
        m = (tt >= 0) & (tt <= 1) & (np.abs(Y - cy) <= tt*(D/2))
    elif shape == "airfoil":           # наклонный эллипс (угол атаки 18°)
        aoa = np.deg2rad(18.0); a = D*0.95; b = D*0.26
        xr = (X - cx)*np.cos(aoa) + (Y - cy)*np.sin(aoa)
        yr = -(X - cx)*np.sin(aoa) + (Y - cy)*np.cos(aoa)
        m = (xr/a)**2 + (yr/b)**2 <= 1
    else:
        m = np.zeros((Ny, Nx), bool)
    m[0, :] = True; m[-1, :] = True    # стенки канала (верх/низ)
    return m

def run_obstacle(shape, Nx=150, Ny=72, D=14, U=0.06, Re=150.0, steps=9600,
                 ramp=1000, frame_every=40):
    cx, cy = Nx//4, Ny//2
    solid_np = solid_mask(Nx, Ny, shape, cx, cy, D)
    solid = cp.asarray(solid_np)
    beta = cm.nu_to_beta(U*D/Re)
    f = E.feq(cp.ones((Ny, Nx), cm.DTYPE), cp.zeros((2, Ny, Nx), cm.DTYPE))
    frames, probe = [], []
    px, py = cx + 3*D, cy
    ones_col = cp.ones((Ny, 1), cm.DTYPE)
    prog = cm.Viz.make_progress(shape)
    for t in range(steps + 1):
        rho, u = E.macros(f)
        fcol = Coll.kbc_collide(f, E.feq(rho, u), beta)
        for i in range(Q):
            fcol[i, solid] = f[OPP[i], solid]               # узловой bounce-back (тело + стенки)
        f = Strm.stream(fcol)
        Uin = U*cm.smoothstep(t/ramp)                       # плавный разгон входа
        uin = cp.zeros((2, Ny, 1), cm.DTYPE); uin[0, :, 0] = Uin
        f[:, 1:-1, 0] = E.feq(ones_col, uin)[:, 1:-1, 0]    # вток (Дирихле)
        f[:, 1:-1, -1] = f[:, 1:-1, -2]                     # отток (нуль-градиент)
        if t % frame_every == 0:
            _, uu = E.macros(f)
            frames.append((t, cm.to_cpu(uu)))
            probe.append(float(uu[1, py, px]))
        prog(t, steps)
    return dict(shape=shape, solid=solid_np, frames=frames, probe=np.array(probe),
                Nx=Nx, Ny=Ny, U=U, Re=Re, D=D, cx=cx, cy=cy)

SHAPES = [("circle", "цилиндр"), ("square", "квадрат"),
          ("triangle", "треугольник"), ("airfoil", "профиль")]
obs = {key: run_obstacle(key) for key, _ in SHAPES}
out_lines = ["Nx=150 Ny=72 D=14 U=0.06 Re=150 steps=9600"]
for key, name in SHAPES:
    p = obs[key]["probe"]; half = p[len(p)//2:]
    out_lines.append(f"{key}: max|u_y(probe)|={np.abs(half).max():.4f}")
    print(f"{name:12}: max|u_y(зонд)|={np.abs(half).max():.4f} (колебания → срыв вихрей)")

# --- гифка-монтаж всех четырёх тел -------------------------------------------------------------
_keys = [k for k, _ in SHAPES]
_titles = {k: n for k, n in SHAPES}
_nf = min(len(obs[k]["frames"]) for k in _keys)
_vmax = max(np.percentile(np.abs(vorticity(obs[k]["frames"][_nf//2][1])), 99.0) for k in _keys)
_ocmap = plt.get_cmap("RdBu_r").copy(); _ocmap.set_bad("#1b1b1b")

def _draw(fig, i):
    for j, k in enumerate(_keys):
        ax = fig.add_subplot(2, 2, j + 1)
        t, u = obs[k]["frames"][i]
        vort = np.where(obs[k]["solid"], np.nan, vorticity(u))
        ax.imshow(vort, origin="lower", cmap=_ocmap, vmin=-_vmax, vmax=_vmax)
        ax.set_title(f"{_titles[k]} · t={t}", fontsize=9)
        ax.set_xticks([]); ax.set_yticks([])
    fig.suptitle("Обтекание тел (KBC, Re≈150): завихренность", fontsize=11)
    fig.tight_layout(rect=(0, 0, 1, 0.96))

cm.render_gif(_nf, _draw, "obstacle_shapes.gif", (11, 6.2), fps=24)

# --- статические панели ------------------------------------------------------------------------
fig, axs = plt.subplots(2, 2, figsize=(12, 5.6), constrained_layout=True)
for ax, k in zip(axs.ravel(), _keys):
    t, u = obs[k]["frames"][-1]
    vort = np.where(obs[k]["solid"], np.nan, vorticity(u))
    ax.imshow(vort, origin="lower", cmap=_ocmap, vmin=-_vmax, vmax=_vmax)
    ax.set_title(f"{_titles[k]} (t={t})", fontsize=10)
    ax.set_xticks([]); ax.set_yticks([])
fig.suptitle("Обтекание тел: поле завихренности в конце прогона", fontweight="bold")
cm.finish(fig, "10_obstacle_shapes.png")

fig, ax = plt.subplots(figsize=(9.5, 3.4))
for k in _keys:
    ax.plot(np.arange(len(obs[k]["probe"])), obs[k]["probe"], label=_titles[k])
ax.set_title("Поперечная скорость в следе (зонд за телом): колебания = срыв вихрей")
ax.set_xlabel("кадр"); ax.set_ylabel("$u_y$ в зонде [lu/ts]"); ax.legend(fontsize=8, ncol=4)
cm.finish(fig, "10b_obstacle_probe.png")

# фон для схемы блочного измельчения (figures_static.py): |ω| цилиндра в конце прогона
rc = obs["circle"]
np.savez(os.path.join(cm.OUTDIR, "obstacle_circle_field.npz"),
         u=rc["frames"][-1][1], solid=rc["solid"],
         Nx=rc["Nx"], Ny=rc["Ny"], cx=rc["cx"], cy=rc["cy"], D=rc["D"])

out_lines.append("PASS  дорожки Кармана на всех четырёх телах")
cm.write_out("obstacle.txt", out_lines)
print("DONE (obstacle)")