# Демо-3: каверна с движущейся крышкой (Re=1000) против эталона Ghia, Ghia & Shin (1982).
# Стенки — bounce-back; крышка — bounce-back с поправкой импульса −2wᵢρ(cᵢ·u_w)/cs².
# Сетка намеренно грубая, BB первого порядка → совпадение качественное.
# Артефакты: figures/09_cavity_ghia.png, anim/cavity_transient.gif, out/cavity.txt|npz.
# Запуск: python cavity_gpu.py
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; OPP = cm.OPP; Q = cm.Q; CS2 = cm.CS2

UP = [2, 5, 6]      # c_iy = +1
DOWN = [4, 7, 8]    # c_iy = -1
RIGHT = [1, 5, 8]   # c_ix = +1
LEFT = [3, 6, 7]    # c_ix = -1

def run_cavity(N=110, Re=1000.0, U=0.05, steps=14000, tol=1e-7, check=500, frame_every=58):
    nu = U*N/Re
    beta = cm.nu_to_beta(nu)
    rho = cp.ones((N, N), cm.DTYPE)
    u = cp.zeros((2, N, N), cm.DTYPE)
    f = E.feq(rho, u)
    W = cm.to_cpu(cm.L.Wa).astype(float)
    prev = None
    frames = []
    prog = cm.Viz.make_progress("cavity")
    t = 0
    for t in range(steps):
        rho, u = E.macros(f)
        if t % frame_every == 0:
            frames.append((t, cm.to_cpu(u)))
        fcol = Coll.kbc_collide(f, E.feq(rho, u), beta)
        f = Strm.stream(fcol)
        # bounce-back: на каждой стенке достраиваем входящие в область направления
        for i in UP:        # низ (y=0)
            f[i, 0, :] = fcol[OPP[i], 0, :]
        for i in RIGHT:     # лево (x=0)
            f[i, :, 0] = fcol[OPP[i], :, 0]
        for i in LEFT:      # право (x=N-1)
            f[i, :, -1] = fcol[OPP[i], :, -1]
        for i in DOWN:      # верх (крышка) + поправка подвижной стенки
            io = OPP[i]
            corr = 2.0*W[io]*(cm.Cx[io]*U)/CS2          # u_w=(U,0): c_y-член равен нулю
            f[i, -1, :] = fcol[io, -1, :] - corr
        if t % check == 0 and t > 0:
            cur = u.copy()
            if prev is not None and float(xp.abs(cur - prev).max()) < tol:
                print(f"\ncavity: сошлось по полю скорости на шаге {t}")
                break
            prev = cur
        prog(t, steps - 1)
    rho, u = E.macros(f)
    return dict(N=N, Re=Re, U=U, u=cm.to_cpu(u), steps=t, frames=frames)

cav = run_cavity()
print(f"Каверна: Re={cav['Re']:.0f}, N={cav['N']}, остановлено на шаге {cav['steps']}")

# Эталон Ghia, Ghia & Shin (1982), Re=1000.
ghia_y = np.array([0.0000, 0.0547, 0.0625, 0.0703, 0.1016, 0.1719, 0.2813, 0.4531,
                   0.5000, 0.6172, 0.7344, 0.8516, 0.9531, 0.9609, 0.9688, 0.9766, 1.0000])
ghia_u = np.array([0.0000, -0.18109, -0.20196, -0.22220, -0.29730, -0.38289, -0.27805,
                   -0.10648, -0.06080, 0.05702, 0.18719, 0.33304, 0.46604, 0.51117,
                   0.57492, 0.65928, 1.0000])
ghia_x = np.array([0.0000, 0.0625, 0.0703, 0.0781, 0.0938, 0.1563, 0.2266, 0.2344,
                   0.5000, 0.8047, 0.8594, 0.9063, 0.9453, 0.9531, 0.9609, 0.9688, 1.0000])
ghia_v = np.array([0.0000, 0.27485, 0.29012, 0.30353, 0.32627, 0.37095, 0.33075, 0.32235,
                   0.02526, -0.31966, -0.42665, -0.51550, -0.39188, -0.33714, -0.27669,
                   -0.21388, 0.0000])

Ncav = cav["N"]
yc = (np.arange(Ncav) + 0.5)/Ncav
xc = (np.arange(Ncav) + 0.5)/Ncav
u_centerline = cav["u"][0, :, Ncav//2]/cav["U"]
v_centerline = cav["u"][1, Ncav//2, :]/cav["U"]

fig, axs = plt.subplots(1, 3, figsize=(14, 4.4))
spd = np.sqrt(cav["u"][0]**2 + cav["u"][1]**2)/cav["U"]
axs[0].imshow(spd, origin="lower", cmap="viridis", extent=[0, 1, 0, 1])
sx = np.linspace(0, 1, Ncav)
axs[0].streamplot(sx, sx, cav["u"][0], cav["u"][1], density=1.1, color="white", linewidth=0.6)
axs[0].set_title("Линии тока и $|\\mathbf{u}|/U$"); axs[0].set_xlabel("$x/L$"); axs[0].set_ylabel("$y/L$")
axs[1].plot(u_centerline, yc, "-", color="#1f6feb", label="KBC")
axs[1].plot(ghia_u, ghia_y, "o", ms=5, mfc="none", color="#d1495b", label="Ghia 1982")
axs[1].set_title("$u_x/U$ по вертикали $x=0.5$"); axs[1].set_xlabel("$u_x/U$")
axs[1].set_ylabel("$y/L$"); axs[1].legend()
axs[2].plot(xc, v_centerline, "-", color="#1f6feb", label="KBC")
axs[2].plot(ghia_x, ghia_v, "o", ms=5, mfc="none", color="#d1495b", label="Ghia 1982")
axs[2].set_title("$u_y/U$ по горизонтали $y=0.5$"); axs[2].set_xlabel("$x/L$")
axs[2].set_ylabel("$u_y/U$"); axs[2].legend()
fig.suptitle("Демо-3: каверна с движущейся крышкой, Re=1000 (KBC vs Ghia)",
             y=1.04, fontweight="bold")
cm.finish(fig, "09_cavity_ghia.png")

u_interp = np.interp(ghia_y, yc, u_centerline)
rms = float(np.sqrt(np.mean((u_interp - ghia_u)**2)))
print(f"Каверна: RMS-отклонение u_x от Ghia ~ {rms:.3f} (грубое BB → допустимо)")

# --- гифка становления течения -----------------------------------------------------------------
_cf = cav["frames"]
_vmax = max(np.sqrt(u[0]**2 + u[1]**2).max() for _, u in _cf)/cav["U"]

def _draw(fig, i):
    t, u = _cf[i]
    ax = fig.add_subplot(111)
    sp = np.sqrt(u[0]**2 + u[1]**2)/cav["U"]
    ax.imshow(sp, origin="lower", cmap="viridis", vmin=0, vmax=_vmax, extent=[0, 1, 0, 1])
    ax.set_title(f"Каверна Re=1000 · t={t} · $|\\mathbf{{u}}|/U$", fontsize=10)
    ax.set_xlabel("$x/L$"); ax.set_ylabel("$y/L$")
    fig.tight_layout()

cm.render_gif(len(_cf), _draw, "cavity_transient.gif", (4.9, 4.5), fps=24)

np.savez(__import__("os").path.join(cm.OUTDIR, "cavity.npz"),
         yc=yc, xc=xc, u_centerline=u_centerline, v_centerline=v_centerline,
         ghia_y=ghia_y, ghia_u=ghia_u, ghia_x=ghia_x, ghia_v=ghia_v, rms=rms,
         Re=cav["Re"], N=cav["N"], steps=cav["steps"])
cm.write_out("cavity.txt", [
    f"N={cav['N']} Re={cav['Re']:.0f} U={cav['U']} stopped_at={cav['steps']}",
    f"rms_ux_vs_Ghia={rms:.4f}",
    "PASS  каверна: качественное совпадение с Ghia (BB первого порядка)",
])
print("DONE (cavity)")