#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
水平卧式流化床 空床流场 —— 批量出图
====================================
读取 fbed_fields.npz, 生成大量切面 / 分量 / 压力 / 分区图。

    python tools/fbed_figures.py
"""
import os
import sys

import numpy as np

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import fbed_simulation as fs

OUT = fs.OUT
BG, FG, MUTED, ACCENT = fs.BG, fs.FG, fs.MUTED, fs.ACCENT
CMAP_FLOW = fs.CMAP_FLOW
XC, YC, ZC = fs.XC, fs.YC, fs.ZC
X0, X1, Y0, Y1, Z0, Z1 = fs.X0, fs.X1, fs.Y0, fs.Y1, fs.Z0, fs.Z1
PPM, DX, SCALE_U = fs.PPM, fs.DX, fs.SCALE_U

D = np.load(os.path.join(OUT, "fbed_fields.npz"))
UX, UY, UZ, RHO, SOLID = D["ux"], D["uy"], D["uz"], D["rho"], D["solid"]
SPEED = np.sqrt(UX ** 2 + UY ** 2 + UZ ** 2) * SCALE_U
P = (RHO - 1.0) / 3.0 * (1.2 * SCALE_U ** 2)
P[SOLID] = np.nan

CMAP_P = plt.get_cmap("RdBu_r")


def xi(v):
    return fs._xi(v)


def yi(v):
    return fs._yi(v)


def zi(v):
    return fs._zi(v)


def get(which, idx):
    if which == "xz":
        return (SPEED[:, idx, :], UX[:, idx, :], UZ[:, idx, :], SOLID[:, idx, :],
                [X0, X1, Z0, Z1], "x (m)", "z (m)")
    if which == "yz":
        return (SPEED[:, :, idx], UY[:, :, idx], UZ[:, :, idx], SOLID[:, :, idx],
                [Y0, Y1, Z0, Z1], "y (m)", "z (m)")
    return (SPEED[idx, :, :], UX[idx, :, :], UY[idx, :, :], SOLID[idx, :, :],
            [X0, X1, Y0, Y1], "x (m)", "y (m)")


def get_scalar(field, which, idx):
    if which == "xz":
        return field[:, idx, :], SOLID[:, idx, :], [X0, X1, Z0, Z1], "x (m)", "z (m)"
    if which == "yz":
        return field[:, :, idx], SOLID[:, :, idx], [Y0, Y1, Z0, Z1], "y (m)", "z (m)"
    return field[idx, :, :], SOLID[idx, :, :], [X0, X1, Y0, Y1], "x (m)", "y (m)"


def panel(ax, which, idx, vmax, stream=True, density=1.1):
    s2, u2, v2, sol2, ext, xl, yl = get(which, idx)
    ax.imshow(np.where(sol2, np.nan, s2), origin="lower", extent=ext, cmap=CMAP_FLOW,
              vmin=0, vmax=vmax, interpolation="bilinear", aspect="equal")
    if stream:
        ny, nx = sol2.shape
        xs = np.linspace(ext[0], ext[1], nx)
        ys = np.linspace(ext[2], ext[3], ny)
        U = np.where(sol2, np.nan, u2)
        V = np.where(sol2, np.nan, v2)
        with np.errstate(invalid="ignore"):
            ax.streamplot(xs, ys, U, V, density=density,
                          color=(0.85, 0.88, 0.92, 0.38), linewidth=0.45,
                          arrowsize=0.5, zorder=5)
    ax.contour(np.linspace(ext[0], ext[1], sol2.shape[1]),
               np.linspace(ext[2], ext[3], sol2.shape[0]),
               sol2.astype(float), levels=[0.5], colors="#adbac7", linewidths=0.7, zorder=6)
    ax.set_xlabel(xl, fontsize=8)
    ax.set_ylabel(yl, fontsize=8)
    ax.set_xlim(ext[0], ext[1])
    ax.set_ylim(ext[2], ext[3])
    ax.tick_params(labelsize=7)


def panel_scalar(ax, field, which, idx, cmap, vmin=None, vmax=None, center=False, stream=True):
    f2, sol2, ext, xl, yl = get_scalar(field, which, idx)
    f2 = np.where(sol2, np.nan, f2)
    if center:
        lim = np.nanpercentile(np.abs(f2), 99) * 1.05
        im = ax.imshow(f2, origin="lower", extent=ext, cmap=cmap, vmin=-lim, vmax=lim,
                       interpolation="bilinear", aspect="equal")
    else:
        im = ax.imshow(f2, origin="lower", extent=ext, cmap=cmap, vmin=vmin, vmax=vmax,
                       interpolation="bilinear", aspect="equal")
    if stream:
        s2, u2, v2, _, _, _, _ = get(which, idx)
        ny, nx = sol2.shape
        xs = np.linspace(ext[0], ext[1], nx)
        ys = np.linspace(ext[2], ext[3], ny)
        with np.errstate(invalid="ignore"):
            ax.streamplot(xs, ys, np.where(sol2, np.nan, u2), np.where(sol2, np.nan, v2),
                          density=1.0, color=(0.85, 0.88, 0.92, 0.3), linewidth=0.4,
                          arrowsize=0.5, zorder=5)
    ax.contour(np.linspace(ext[0], ext[1], sol2.shape[1]),
               np.linspace(ext[2], ext[3], sol2.shape[0]),
               sol2.astype(float), levels=[0.5], colors="#adbac7", linewidths=0.7, zorder=6)
    ax.set_xlabel(xl, fontsize=8)
    ax.set_ylabel(yl, fontsize=8)
    ax.set_xlim(ext[0], ext[1])
    ax.set_ylim(ext[2], ext[3])
    ax.tick_params(labelsize=7)
    return im


def montage(specs, ncols, suptitle, name, vmax, cbar_label="|U| (m/s)",
            figw=17.0, panel_h=3.0, stream=True, density=1.1):
    nrows = (len(specs) + ncols - 1) // ncols
    fig, axes = plt.subplots(nrows, ncols, figsize=(figw, panel_h * nrows + 1.0),
                             squeeze=False)
    for ax in axes.ravel()[len(specs):]:
        ax.axis("off")
    for ax, (spec, title) in zip(axes.ravel(), specs):
        which, idx = spec
        panel(ax, which, idx, vmax, stream=stream, density=density)
        ax.set_title(title, fontsize=9.5, pad=4)
    sm = plt.cm.ScalarMappable(cmap=CMAP_FLOW, norm=plt.Normalize(0, vmax))
    cb = fig.colorbar(sm, ax=axes.ravel().tolist(), fraction=0.012, pad=0.01)
    cb.set_label(cbar_label, fontsize=10)
    cb.outline.set_edgecolor("#30363d")
    fig.suptitle(suptitle, fontsize=15, y=0.995)
    fig.savefig(os.path.join(OUT, name), dpi=500, bbox_inches="tight")
    plt.close(fig)
    print("  ->", name)


def fig_xy_plenum():
    specs = [(("xy", zi(z)), f"z = {z:+.2f} m") for z in (-0.42, -0.32, -0.22, -0.12, -0.02)]
    montage(specs, 3, "风箱层水平切面 |U|（4 段侧进风分布）", "fbed-montage-plenum.png", 6.0)


def fig_xy_plate():
    specs = [(("xy", zi(z)), f"z = {z:.2f} m") for z in (0.04, 0.08, 0.13, 0.20, 0.30, 0.45)]
    montage(specs, 3, "布风板近区水平切面 |U|（射流合并过程）", "fbed-montage-nearplate.png", 3.0)


def fig_xy_free():
    specs = [(("xy", zi(z)), f"z = {z:.2f} m") for z in (0.60, 0.80, 1.00, 1.20, 1.45, 1.75)]
    montage(specs, 3, "沸腾段 / 渐扩 / 排风 水平切面 |U|", "fbed-montage-freeboard.png", 4.0)


def fig_xz():
    specs = [(("xz", yi(y)), f"y = {y:+.2f} m") for y in (-0.26, -0.13, 0.0, 0.13, 0.26)]
    montage(specs, 3, "垂直切面 x-z（不同宽度位置）|U|", "fbed-montage-xz.png", 5.0, figw=16.0)


def fig_yz():
    specs = [(("yz", xi(x)), f"x = {x:.2f} m") for x in (0.15, 0.45, 0.75, 1.05, 1.35, 1.65, 1.95, 2.25)]
    montage(specs, 4, "垂直切面 y-z（沿长度）|U|", "fbed-montage-yz.png", 5.0, figw=18.0, panel_h=3.4)


def fig_components():
    fig, axes = plt.subplots(2, 2, figsize=(15.0, 12.0))
    panel(axes[0, 0], "xz", yi(0.0), 4.0)
    axes[0, 0].set_title("(a) 速度大小 |U|", fontsize=12, color=ACCENT)
    for ax, comp, title, cmap in [
            (axes[0, 1], UX, "(b) 水平分量 ux（孔板 5° 倾斜 + 向排风汇流）", "coolwarm"),
            (axes[1, 0], UY, "(c) 横向分量 uy（侧进风 + 渐扩）", "coolwarm"),
            (axes[1, 1], UZ, "(d) 垂直分量 uz（上升主流）", "coolwarm")]:
        f2, sol2, ext, xl, yl = get_scalar(comp * SCALE_U, "xz", yi(0.0))
        lim = np.nanpercentile(np.abs(np.where(sol2, np.nan, f2)), 99) * 1.05
        im = ax.imshow(np.where(sol2, np.nan, f2), origin="lower", extent=ext, cmap=cmap,
                       vmin=-lim, vmax=lim, interpolation="bilinear", aspect="equal")
        ax.contour(np.linspace(ext[0], ext[1], sol2.shape[1]),
                   np.linspace(ext[2], ext[3], sol2.shape[0]),
                   sol2.astype(float), levels=[0.5], colors="#adbac7", linewidths=0.7)
        ax.set_title(title, fontsize=12)
        ax.set_xlabel("x (m)", fontsize=9)
        ax.set_ylabel("z (m)", fontsize=9)
        cb = fig.colorbar(im, ax=ax, fraction=0.036, pad=0.02)
        cb.set_label("m/s", fontsize=9)
        cb.outline.set_edgecolor("#30363d")
    fig.suptitle("速度分量分解（垂直切面 y=0）", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(os.path.join(OUT, "fbed-components.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-components.png")


def fig_pressure():
    fig, axes = plt.subplots(2, 2, figsize=(15.0, 12.0))
    for ax, (which, idx, title) in zip(axes.ravel(), [
            ("xz", yi(0.0), "(a) 垂直切面 y=0 —— 全系统压差"),
            ("yz", xi(1.2), "(b) 垂直切面 x=1.2"),
            ("xy", zi(0.15), "(c) 布风板近区 z=0.15"),
            ("xy", zi(0.90), "(d) 沸腾段 z=0.90")]):
        im = panel_scalar(ax, P, which, idx, CMAP_P, center=True, stream=False)
        ax.set_title(title, fontsize=12)
        cb = fig.colorbar(im, ax=ax, fraction=0.036, pad=0.02)
        cb.set_label("静压 (Pa)", fontsize=9)
        cb.outline.set_edgecolor("#30363d")
    fig.suptitle("静压场（相对出口, 单位 Pa）", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(os.path.join(OUT, "fbed-pressure.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-pressure.png")


def fig_sections():
    fig, axes = plt.subplots(1, 4, figsize=(17.0, 6.0))
    for k, ax in enumerate(axes):
        xc = 0.6 * (k + 0.5)
        panel(ax, "yz", xi(xc), 5.0)
        ax.set_title(f"第 {k+1} 段  x={xc:.1f} m", fontsize=11)
    fig.suptitle("4 个进风区段中心垂直切面 y-z（单段风量 2500 m³/h）", fontsize=14, y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(os.path.join(OUT, "fbed-sections.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-sections.png")


def fig_hero_stream():
    fig, axes = plt.subplots(1, 2, figsize=(18.0, 8.0))
    panel(axes[0], "xz", yi(0.0), 4.0, density=1.8)
    axes[0].set_title("垂直切面 y=0 —— 风箱/布风板/沸腾段/排风", fontsize=12)
    panel(axes[1], "yz", xi(1.2), 4.0, density=1.8)
    axes[1].set_title("垂直切面 x=1.2 —— 侧向渐扩", fontsize=12)
    fig.suptitle("空床流场流线（GPU D3Q19 LBM）", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig(os.path.join(OUT, "fbed-streamlines.png"), dpi=560)
    plt.close(fig)
    print("  -> fbed-streamlines.png")


def fig_3d():
    from mpl_toolkits.mplot3d import Axes3D  # noqa
    fig = plt.figure(figsize=(13.0, 10.0))
    ax = fig.add_subplot(111, projection="3d")
    norm = plt.Normalize(0, 4.0)

    def surf_y(iy):
        f = SPEED[:, iy, :]
        m = SOLID[:, iy, :]
        X, Z = np.meshgrid(XC, ZC)
        Y = np.full_like(X, YC[iy])
        fc = CMAP_FLOW(norm(np.where(m, np.nan, f)))
        fc[m, 3] = 0.0
        ax.plot_surface(X, Y, Z, facecolors=fc, shade=False, rstride=1, cstride=1,
                        linewidth=0, antialiased=False)

    def surf_x(ix):
        f = SPEED[:, :, ix]
        m = SOLID[:, :, ix]
        Y, Z = np.meshgrid(YC, ZC)
        X = np.full_like(Y, XC[ix])
        fc = CMAP_FLOW(norm(np.where(m, np.nan, f)))
        fc[m, 3] = 0.0
        ax.plot_surface(X, Y, Z, facecolors=fc, shade=False, rstride=1, cstride=1,
                        linewidth=0, antialiased=False)

    def surf_z(iz):
        f = SPEED[iz, :, :]
        m = SOLID[iz, :, :]
        X, Y = np.meshgrid(XC, YC)
        Z = np.full_like(X, ZC[iz])
        fc = CMAP_FLOW(norm(np.where(m, np.nan, f)))
        fc[m, 3] = 0.0
        ax.plot_surface(X, Y, Z, facecolors=fc, shade=False, rstride=1, cstride=1,
                        linewidth=0, antialiased=False)

    surf_y(yi(0.0))
    surf_x(xi(1.2))
    surf_z(zi(0.9))
    ax.set_xlabel("x (m)")
    ax.set_ylabel("y (m)")
    ax.set_zlabel("z (m)")
    ax.set_title("三正交切面三维剖切（|U|）", fontsize=13)
    ax.view_init(elev=22, azim=-58)
    sm = plt.cm.ScalarMappable(cmap=CMAP_FLOW, norm=norm)
    cb = fig.colorbar(sm, ax=ax, fraction=0.03, pad=0.08, shrink=0.6)
    cb.set_label("|U| (m/s)", fontsize=10)
    cb.outline.set_edgecolor("#30363d")
    fig.savefig(os.path.join(OUT, "fbed-3d-cutaway.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-3d-cutaway.png")


def main():
    print("[figs] loading", os.path.join(OUT, "fbed_fields.npz"))
    fig_xy_plenum()
    fig_xy_plate()
    fig_xy_free()
    fig_xz()
    fig_yz()
    fig_components()
    fig_pressure()
    fig_sections()
    fig_hero_stream()
    fig_3d()
    print("[figs] done")


if __name__ == "__main__":
    main()