# Демо-2: двойной сдвиговый слой (Minion–Brown) на недоразрешённой сетке — устойчивость
# BGK (poly-равновесие) против KBC (entropic + γ): BGK разрушается, KBC держит (implicit LES).
# Артефакты: figures/08_shear_stability.png, anim/shear_bgk_vs_kbc.gif, out/shear.txt.
# Запуск: python shear_gpu.py
import math
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

def init_shear(N, U0=0.04, delta=0.05, kappa=80.0):
    xs = cp.arange(N, dtype=cm.DTYPE)/N
    X, Yc = cp.meshgrid(xs, xs)
    ux = cp.where(Yc <= 0.5, U0*cp.tanh(kappa*(Yc - 0.25)), U0*cp.tanh(kappa*(0.75 - Yc)))
    uy = delta*U0*cp.sin(2.0*math.pi*(X + 0.25))
    return cm.B.pack((ux.astype(cm.DTYPE), uy.astype(cm.DTYPE)))

def run_shear(scheme, N=128, U0=0.04, nu=1.7e-4, steps=8000, rec=100, frame_every=33):
    u0 = init_shear(N, U0=U0)
    rho = cp.ones((N, N), cm.DTYPE)
    beta = cm.nu_to_beta(nu)
    f = E.feq(rho, u0) if scheme == "kbc" else cm.feq_poly(rho, u0)
    ts, umax, frames = [], [], []
    blew_up_at = None
    last_u = cm.to_cpu(u0)
    prog = cm.Viz.make_progress(scheme)
    for t in range(steps + 1):
        rho, u = E.macros(f)
        um = float(xp.abs(u).max())
        if not math.isfinite(um) or um > 0.5:
            blew_up_at = t
            print(f"\n{scheme}: разрушился на шаге {t}")
            break
        if t % frame_every == 0:
            frames.append((t, cm.to_cpu(u)))
        if scheme == "kbc":
            f = Coll.kbc_collide(f, E.feq(rho, u), beta)
        else:
            f = cm.bgk_collide(f, cm.feq_poly(rho, u), beta)
        f = Strm.stream(f)
        last_u = cm.to_cpu(u)
        if t % rec == 0:
            ts.append(t); umax.append(um)
        prog(t, steps)
    gamma = None
    if scheme == "kbc":                                # поле γ на финальном состоянии — для панели
        r, u = E.macros(f)
        gamma = cm.to_cpu(cm.kbc_gamma(f, E.feq(r, u), beta))
    return dict(scheme=scheme, ts=np.array(ts), umax=np.array(umax), frames=frames,
                u=last_u, gamma=gamma, blew_up_at=blew_up_at, N=N, U0=U0, nu=nu)

Re_shear = 0.04*128/1.7e-4
print(f"Сдвиговый слой: Re ~ {Re_shear:.0f}, β ~ {cm.nu_to_beta(1.7e-4):.4f}")
res_bgk = run_shear("bgk")
res_kbc = run_shear("kbc")
print("BGK:", (f"разрушился на t={res_bgk['blew_up_at']}" if res_bgk["blew_up_at"] else "устойчив"))
print("KBC:", (f"разрушился на t={res_kbc['blew_up_at']}" if res_kbc["blew_up_at"] else "устойчив"))
assert res_kbc["blew_up_at"] is None, "KBC должен оставаться устойчивым"

# --- панель: ω(KBC), поле γ−2, max|u|(t) -------------------------------------------------------
fig, axs = plt.subplots(1, 3, figsize=(13.5, 4.0))
om_k = vorticity(res_kbc["u"])
im0 = axs[0].imshow(om_k, origin="lower", cmap="RdBu_r")
axs[0].set_title("KBC: завихренность (устойчиво)")
axs[0].set_xlabel("$x$ [lu]"); axs[0].set_ylabel("$y$ [lu]")
plt.colorbar(im0, ax=axs[0], fraction=0.046)
im1 = axs[1].imshow(res_kbc["gamma"] - 2.0, origin="lower", cmap="magma")
axs[1].set_title("Поле стабилизатора $\\gamma-2$")
axs[1].set_xlabel("$x$ [lu]"); axs[1].set_ylabel("$y$ [lu]")
plt.colorbar(im1, ax=axs[1], fraction=0.046)
axs[2].plot(res_kbc["ts"], res_kbc["umax"], color="#2e7d32", label="KBC")
axs[2].plot(res_bgk["ts"], res_bgk["umax"], color="#d1495b", label="BGK")
if res_bgk["blew_up_at"]:
    axs[2].axvline(res_bgk["blew_up_at"], color="#d1495b", ls="--", alpha=0.6)
axs[2].set_title("$\\max|\\mathbf{u}|$ во времени")
axs[2].set_xlabel("$t$ [ts]"); axs[2].set_ylabel("$\\max|\\mathbf{u}|$ [lu/ts]"); axs[2].legend()
fig.suptitle(f"Демо-2: сдвиговый слой, Re~{Re_shear:.0f} (BGK разрушается, KBC держит)",
             y=1.04, fontweight="bold")
cm.finish(fig, "08_shear_stability.png")

# --- гифка: BGK (слева, замирает с пометкой) vs KBC (справа) -----------------------------------
_kf, _bf = res_kbc["frames"], res_bgk["frames"]
_om0s = np.percentile(np.abs(vorticity(_kf[len(_kf)//2][1])), 99.5)

def _draw(fig, i):
    ax1 = fig.add_subplot(1, 2, 1)
    ax2 = fig.add_subplot(1, 2, 2)
    dead = i >= len(_bf)
    tb, ub = _bf[min(i, len(_bf) - 1)]
    ax1.imshow(vorticity(ub), origin="lower", cmap="RdBu_r", vmin=-_om0s, vmax=_om0s)
    ax1.set_title("BGK — разрушился" if dead else f"BGK · t={tb}",
                  color="red" if dead else "black", fontsize=10)
    tk, uk = _kf[i]
    ax2.imshow(vorticity(uk), origin="lower", cmap="RdBu_r", vmin=-_om0s, vmax=_om0s)
    ax2.set_title(f"KBC · t={tk}", fontsize=10)
    for a in (ax1, ax2):
        a.set_xticks([]); a.set_yticks([])
    fig.tight_layout()

cm.render_gif(len(_kf), _draw, "shear_bgk_vs_kbc.gif", (8.4, 4.3), fps=24)

cm.write_out("shear.txt", [
    f"N=128 U0=0.04 nu=1.7e-4 Re~{Re_shear:.0f} steps=8000",
    f"BGK_blew_up_at={res_bgk['blew_up_at']}",
    f"KBC_stable={res_kbc['blew_up_at'] is None}",
    "PASS  KBC устойчив там, где BGK разрушается",
])
print("DONE (shear)")