# Демо-1: вихрь Тейлора–Грина (валидация вязкости и H-теоремы) на компонентах solver_2x_sdf.
# Аналитика: E(t) = E0·exp(−4νk²t). Меряем ν по наклону ln E(t) — должен совпасть с заданным.
# Артефакты: figures/07_taylor_green.png, figures/07b_tgv_decay_slices.png, anim/tgv_decay.gif,
# out/tgv.txt. Запуск: python tgv_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 run_tgv(N=96, U0=0.04, nu=0.0125, steps=10000, rec=50, n_snap=10, frame_every=42):
    k = 2.0*math.pi/N
    xs = cp.arange(N, dtype=cm.DTYPE)
    X, Yc = cp.meshgrid(xs, xs)                       # X — ось 1, Yc — ось 0
    ux = -U0*cp.cos(k*X)*cp.sin(k*Yc)
    uy = U0*cp.sin(k*X)*cp.cos(k*Yc)
    rho = cp.ones((N, N), cm.DTYPE)
    f = E.feq(rho, cm.B.pack((ux.astype(cm.DTYPE), uy.astype(cm.DTYPE))))
    beta = cm.nu_to_beta(nu)
    snap_at = sorted(set(int(round(v)) for v in np.linspace(0, steps, n_snap)))
    snaps, frames, ts, KE, Hs = [], [], [], [], []
    prog = cm.Viz.make_progress("tgv")
    for t in range(steps + 1):
        rho, u = E.macros(f)
        f = Coll.kbc_collide(f, E.feq(rho, u), beta)
        f = Strm.stream(f)                            # периодический перенос — ровно случай TGV
        if (t in snap_at) or (t % rec == 0) or (t % frame_every == 0):
            _, uu = E.macros(f)
            if t in snap_at:
                snaps.append((t, cm.to_cpu(uu)))
            if t % frame_every == 0:
                frames.append((t, cm.to_cpu(uu)))
            if t % rec == 0:
                ts.append(t)
                KE.append(float((0.5*(uu[0]**2 + uu[1]**2)).mean()))
                Hs.append(float(cm.H_function(f).mean()))
        prog(t, steps)
    return dict(N=N, k=k, nu=nu, U0=U0, ts=np.array(ts), KE=np.array(KE),
                H=np.array(Hs), u=cm.to_cpu(E.macros(f)[1]), snaps=snaps, frames=frames)

tgv = run_tgv()
mask = tgv["KE"] > 0
slope = np.polyfit(tgv["ts"][mask], np.log(tgv["KE"][mask]), 1)[0]
nu_meas = -slope/(4.0*tgv["k"]**2)
rel_err = abs(nu_meas - tgv["nu"])/tgv["nu"]
print(f"TGV: ν задано={tgv['nu']:.5f}, ν измерено={nu_meas:.5f}, отн. ошибка={rel_err*100:.2f}%")
assert rel_err < 0.06, "TGV: измеренная вязкость отклонилась слишком сильно"
dH = np.diff(tgv["H"])
print(f"TGV: H-функция монотонно невозрастает: max(ΔH)={dH.max():.2e}")

# --- статическая панель: ω, затухание энергии, H(t) -------------------------------------------
fig, axs = plt.subplots(1, 3, figsize=(13.5, 4.0))
om = vorticity(tgv["u"])
im = axs[0].imshow(om, origin="lower", cmap="RdBu_r")
axs[0].set_title("Завихренность $\\omega$ (TGV, срез)")
axs[0].set_xlabel("$x$ [lu]"); axs[0].set_ylabel("$y$ [lu]")
plt.colorbar(im, ax=axs[0], fraction=0.046)
axs[1].semilogy(tgv["ts"], tgv["KE"], "o", ms=3, label="KBC (измерено)")
axs[1].semilogy(tgv["ts"], tgv["KE"][0]*np.exp(-4*tgv["nu"]*tgv["k"]**2*tgv["ts"]),
                "-", color="#d1495b", label="$E_0 e^{-4\\nu k^2 t}$ (аналитика)")
axs[1].set_title("Затухание кин. энергии"); axs[1].set_xlabel("$t$ [ts]")
axs[1].set_ylabel("$E(t)$"); axs[1].legend()
axs[2].plot(tgv["ts"], tgv["H"], color="#2e7d32")
axs[2].set_title("H-функция (энтропия) во времени")
axs[2].set_xlabel("$t$ [ts]"); axs[2].set_ylabel("$\\langle H\\rangle$")
fig.suptitle("Демо-1: вихрь Тейлора–Грина (KBC, D2Q9)", y=1.04, fontweight="bold")
cm.finish(fig, "07_taylor_green.png")

# --- 10 срезов распада с единой шкалой цвета ---------------------------------------------------
snaps = tgv["snaps"]
om0 = np.abs(vorticity(snaps[0][1])).max()
KE0 = 0.5*np.mean(snaps[0][1][0]**2 + snaps[0][1][1]**2)
fig, axs = plt.subplots(2, 5, figsize=(15, 6.3), constrained_layout=True)
im = None
for ax, (tt, uu) in zip(axs.ravel(), snaps):
    im = ax.imshow(vorticity(uu), origin="lower", cmap="RdBu_r", vmin=-om0, vmax=om0)
    KEf = 0.5*np.mean(uu[0]**2 + uu[1]**2)/KE0
    ax.set_title(f"t={tt} · $E/E_0$={KEf:.2f}", fontsize=10)
    ax.set_xticks([]); ax.set_yticks([])
fig.colorbar(im, ax=axs, fraction=0.025, pad=0.01, label="завихренность $\\omega$ [1/ts]")
fig.suptitle("Распад вихря Тейлора–Грина: завихренность на равных интервалах времени "
             "(единая шкала цвета)", fontweight="bold")
cm.finish(fig, "07b_tgv_decay_slices.png")

# --- гифка затухания ---------------------------------------------------------------------------
_fr = tgv["frames"]
_om0 = np.abs(vorticity(_fr[0][1])).max()
_KE0 = 0.5*np.mean(_fr[0][1][0]**2 + _fr[0][1][1]**2)

def _draw(fig, i):
    t, u = _fr[i]
    ax = fig.add_subplot(111)
    ax.imshow(vorticity(u), origin="lower", cmap="RdBu_r", vmin=-_om0, vmax=_om0)
    ke = 0.5*np.mean(u[0]**2 + u[1]**2)/_KE0
    ax.set_title(f"Тейлор–Грин · t={t} · $E/E_0$={ke:.2f}", fontsize=10)
    ax.set_xticks([]); ax.set_yticks([])
    fig.tight_layout()

cm.render_gif(len(_fr), _draw, "tgv_decay.gif", (4.6, 4.4), fps=24)

cm.write_out("tgv.txt", [
    f"N={tgv['N']} U0={tgv['U0']} steps=10000",
    f"nu_given={tgv['nu']:.5f} nu_measured={nu_meas:.5f} rel_err_pct={rel_err*100:.2f}",
    f"H_monotone_max_dH={dH.max():.3e}",
    "PASS  TGV: вязкость по затуханию энергии и H-теорема",
])
print("DONE (tgv)")