#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
水平卧式流化床 空床流场 —— 分析图 (几何/矢量/压差/流量/均匀性/衰减/倾斜/分段放大)
读取 fbed_fields.npz。

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

import numpy as np

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

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
DX, SCALE_U, BED_L, BED_W, N_SEC, SEC_L = fs.DX, fs.SCALE_U, fs.BED_L, fs.BED_W, fs.N_SEC, fs.SEC_L
PPM = fs.PPM
NZ, NY, NX = fs.NZ, fs.NY, fs.NX

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"]
PLATE = D["plate"]
FLUID = ~SOLID
UX, UY, UZ = UX * SCALE_U, UY * SCALE_U, UZ * SCALE_U
SPEED = np.sqrt(UX ** 2 + UY ** 2 + UZ ** 2)
P = (RHO - 1.0) / 3.0 * (1.2 * SCALE_U ** 2)
CELL_A = DX * DX
IZP = fs._zi(0.0)


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


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


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


def quiver_slice(ax, which, idx, step=3, scale=None, title="", vmax=None):
    if which == "xz":
        ext = [X0, X1, Z0, Z1]; xl, yl = "x (m)", "z (m)"
        U, V, M = UX[:, idx, :], UZ[:, idx, :], SPEED[:, idx, :]
        m = SOLID[:, idx, :]
        xs, ys = XC, ZC
    elif which == "yz":
        ext = [Y0, Y1, Z0, Z1]; xl, yl = "y (m)", "z (m)"
        U, V, M = UY[:, :, idx], UZ[:, :, idx], SPEED[:, :, idx]
        m = SOLID[:, :, idx]
        xs, ys = YC, ZC
    else:
        ext = [X0, X1, Y0, Y1]; xl, yl = "x (m)", "y (m)"
        U, V, M = UX[idx, :, :], UY[idx, :, :], SPEED[idx, :, :]
        m = SOLID[idx, :, :]
        xs, ys = XC, YC
    U = np.where(m, np.nan, U); V = np.where(m, np.nan, V); M = np.where(m, np.nan, M)
    ax.imshow(np.where(m, np.nan, M), origin="lower", extent=ext, cmap=CMAP_FLOW,
              vmin=0, vmax=vmax or np.nanpercentile(M, 99), interpolation="bilinear",
              aspect="equal", alpha=0.85)
    Xs, Ys = np.meshgrid(xs[::step], ys[::step])
    Us = U[::step, ::step]; Vs = V[::step, ::step]; Ms = M[::step, ::step]
    if scale is None:
        scale = 6.0 * max(ext[1] - ext[0], ext[3] - ext[2]) / max(np.nanmax(np.hypot(Us, Vs)), 1e-6)
    ax.quiver(Xs, Ys, Us, Vs, Ms, cmap="cool", scale=scale, width=0.0022,
              clim=(0, np.nanpercentile(M, 99)))
    ax.contour(np.linspace(ext[0], ext[1], m.shape[1]),
               np.linspace(ext[2], ext[3], m.shape[0]),
               m.astype(float), levels=[0.5], colors="#adbac7", linewidths=0.8)
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3])
    ax.set_xlabel(xl, fontsize=9); ax.set_ylabel(yl, fontsize=9)
    ax.tick_params(labelsize=8)
    if title:
        ax.set_title(title, fontsize=12)


# --------------------------------------------------------------------------
def fig_geometry():
    fig, axes = plt.subplots(1, 2, figsize=(17.0, 7.6))
    for ax, which, idx in [(axes[0], "xz", yi(0.0)), (axes[1], "yz", xi(1.2))]:
        if which == "xz":
            ext = [X0, X1, Z0, Z1]; sol = SOLID[:, idx, :]
            xl, yl = "x (m)", "z (m)"
        else:
            ext = [Y0, Y1, Z0, Z1]; sol = SOLID[:, :, idx]
            xl, yl = "y (m)", "z (m)"
        ax.imshow(np.where(sol, 0.55, 0.08), origin="lower", extent=ext, cmap="gray",
                  vmin=0, vmax=1, interpolation="nearest", aspect="equal")
        ax.contour(np.linspace(ext[0], ext[1], sol.shape[1]),
                   np.linspace(ext[2], ext[3], sol.shape[0]),
                   sol.astype(float), levels=[0.5], colors="#adbac7", linewidths=1.0)
        ax.set_xlabel(xl, fontsize=10); ax.set_ylabel(yl, fontsize=10)
    # 标注 (在 x-z 图)
    ax = axes[0]
    ax.annotate("4 x DN300 侧进风\n(平行床板, 2500 m³/h/段)", xy=(0.05, -0.225), xytext=(-0.35, 0.9),
                fontsize=10, color="#7ee787",
                arrowprops=dict(arrowstyle="->", color="#7ee787", lw=1.6))
    ax.annotate("布风板 2mm孔 5.6% 5°倾斜", xy=(1.2, 0.0), xytext=(1.55, -0.35),
                fontsize=10, color="#58a6ff",
                arrowprops=dict(arrowstyle="->", color="#58a6ff", lw=1.4))
    ax.annotate("顶部仅 DN800 排风\n(无进风)", xy=(1.2, 2.1), xytext=(1.45, 1.55),
                fontsize=10.5, color="#ff7b72",
                arrowprops=dict(arrowstyle="->", color="#ff7b72", lw=1.6))
    ax.set_title("几何 / 边界示意  (x-z, y=0)", fontsize=12)
    axes[1].set_title("几何 / 边界示意  (y-z, x=1.2)", fontsize=12)
    fig.suptitle("空床几何与边界条件：底部 4 段侧进风，顶部仅排风", fontsize=15, y=0.98)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(os.path.join(OUT, "fbed-geometry.png"), dpi=540)
    plt.close(fig)
    print("  -> fbed-geometry.png")


def fig_vectors():
    fig, axes = plt.subplots(2, 2, figsize=(16.0, 12.5))
    quiver_slice(axes[0, 0], "xz", yi(0.0), step=3, title="(a) 垂直切面 y=0 速度矢量")
    quiver_slice(axes[0, 1], "yz", xi(1.2), step=3, title="(b) 垂直切面 x=1.2 速度矢量")
    quiver_slice(axes[1, 0], "xy", zi(0.15), step=3, title="(c) 布风板近区 z=0.15 速度矢量")
    quiver_slice(axes[1, 1], "xy", zi(0.90), step=3, title="(d) 沸腾段 z=0.90 速度矢量")
    fig.suptitle("空床流场速度矢量图（箭头方向=流向, 颜色=速度大小 m/s）", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(os.path.join(OUT, "fbed-vectors.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-vectors.png")


def fig_tilt():
    fig = plt.figure(figsize=(16.0, 10.5))
    gs = fig.add_gridspec(2, 2)
    ax = fig.add_subplot(gs[0, 0])
    iz = zi(0.05)
    quiver_slice(ax, "xy", iz, step=4, vmax=2.5,
                 title="(a) 布风板正上方 z=0.05 矢量（放大看 +x 倾斜）")
    ax.set_xlim(0.6, 1.8); ax.set_ylim(-0.30, 0.30)

    ax = fig.add_subplot(gs[0, 1])
    zs, ang, um = [], [], []
    for z in np.arange(0.03, 0.60, 0.02):
        iz = zi(z)
        m = FLUID[iz] & (np.abs(YC)[:, None] < BED_W / 2)
        if m.sum() < 10:
            continue
        ux_m, uz_m = UX[iz][m].mean(), UZ[iz][m].mean()
        zs.append(z); ang.append(np.degrees(np.arctan2(ux_m, uz_m))); um.append(ux_m)
    ax.plot(zs, ang, "-o", color=ACCENT, ms=4, label="实际偏角")
    ax.axhline(5.0, color="#ff7b72", ls="--", lw=1.2, label="理论 +5°")
    ax.axhline(0.0, color=MUTED, lw=0.8)
    ax.set_xlabel("高度 z (m)", fontsize=10); ax.set_ylabel("流向偏角 (°)", fontsize=10)
    ax.set_title("(b) 近板流向偏角随高度衰减", fontsize=12)
    ax.legend(fontsize=9); ax.grid(alpha=0.2)

    ax = fig.add_subplot(gs[1, 0])
    zs2, uxs, uzs = [], [], []
    for z in np.arange(0.03, 0.60, 0.02):
        iz = zi(z)
        m = FLUID[iz] & (np.abs(YC)[:, None] < BED_W / 2)
        if m.sum() < 10:
            continue
        zs2.append(z); uxs.append(UX[iz][m].mean()); uzs.append(UZ[iz][m].mean())
    ax.plot(zs2, uzs, "-o", ms=4, color="#39d0d8", label="垂直分量 <uz>")
    ax.plot(zs2, uxs, "-s", ms=4, color="#ff7b72", label="水平分量 <ux>（5° 驱动）")
    ax.set_xlabel("高度 z (m)", fontsize=10); ax.set_ylabel("面平均速度 (m/s)", fontsize=10)
    ax.set_title("(c) 近板面平均速度分量", fontsize=12)
    ax.legend(fontsize=9); ax.grid(alpha=0.2)

    ax = fig.add_subplot(gs[1, 1])
    xline = np.linspace(0.02, BED_L - 0.02, 60)
    iz = zi(0.06)
    prof = []
    for xv in xline:
        i = xi(xv)
        m = FLUID[iz, :, i]
        prof.append(UX[iz, :, i][m].mean() if m.sum() else np.nan)
    ax.plot(xline, prof, color="#ff7b72", lw=1.6)
    ax.axhline(0, color=MUTED, lw=0.8)
    ax.set_xlabel("x (m)", fontsize=10); ax.set_ylabel("<ux> (m/s)", fontsize=10)
    ax.set_title("(d) 近板水平速度沿长度分布（整体 +x 偏置）", fontsize=12)
    ax.grid(alpha=0.2)
    fig.suptitle("布风板 5° 倾斜对近板风向的影响", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(os.path.join(OUT, "fbed-tilt.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-tilt.png")


def fig_uniformity():
    fig, axes = plt.subplots(1, 3, figsize=(17.0, 5.4))
    qs = []
    for k in range(N_SEC):
        ix0, ix1 = xi(0.6 * k), xi(0.6 * (k + 1))
        q = UZ[IZP, :, ix0:ix1][FLUID[IZP, :, ix0:ix1]].sum() * CELL_A * 3600.0
        qs.append(q)
    ax = axes[0]
    bars = ax.bar([f"第{k+1}段" for k in range(N_SEC)], qs, color=ACCENT, alpha=0.85)
    ax.axhline(2500, color="#ff7b72", ls="--", lw=1.4, label="目标 2500 m³/h")
    for b, q in zip(bars, qs):
        ax.text(b.get_x() + b.get_width() / 2, q + 40, f"{q:.0f}", ha="center", fontsize=9)
    ax.set_ylabel("布风板风量 (m³/h)", fontsize=10); ax.legend(fontsize=9)
    ax.set_title("(a) 各区段风量", fontsize=12); ax.grid(alpha=0.2, axis="y")

    ax = axes[1]
    prof = []
    xline = np.linspace(0.02, BED_L - 0.02, 120)
    for xv in xline:
        i = xi(xv)
        m = FLUID[IZP + 3, :, i]
        prof.append(UZ[IZP + 3, :, i][m].mean() if m.sum() else np.nan)
    ax.plot(xline, prof, color="#39d0d8", lw=1.8)
    ax.axhline(2500.0 / (BED_L * BED_W) / 3600.0 * 3600.0, color=MUTED, ls=":", lw=1)
    for k in range(1, N_SEC):
        ax.axvline(0.6 * k, color="#30363d", lw=0.8)
    ax.set_xlabel("x (m)", fontsize=10); ax.set_ylabel("近板 <uz> (m/s)", fontsize=10)
    ax.set_title("(b) 近板上升速度沿长度分布（均匀性）", fontsize=12); ax.grid(alpha=0.2)

    ax = axes[2]
    yline = np.linspace(-BED_W / 2 + 0.02, BED_W / 2 - 0.02, 60)
    profy = []
    for yv in yline:
        j = yi(yv)
        m = FLUID[IZP + 3, j, :]
        profy.append(UZ[IZP + 3, j, :][m].mean() if m.sum() else np.nan)
    ax.plot(yline, profy, color="#7ee787", lw=1.8)
    ax.set_xlabel("y (m)", fontsize=10); ax.set_ylabel("近板 <uz> (m/s)", fontsize=10)
    ax.set_title("(c) 近板上升速度沿宽度分布（均匀性）", fontsize=12); ax.grid(alpha=0.2)
    fig.suptitle("布风板流量均匀性", fontsize=15, y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(os.path.join(OUT, "fbed-uniformity.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-uniformity.png")


def fig_decay():
    fig, axes = plt.subplots(1, 2, figsize=(15.0, 7.0))
    zs, qz, uz_avg, spd_avg, area = [], [], [], [], []
    for iz in range(NZ):
        m = FLUID[iz]
        if m.sum() < 20:
            continue
        zs.append(ZC[iz])
        qz.append(UZ[iz][m].sum() * CELL_A * 3600.0)
        uz_avg.append(UZ[iz][m].mean())
        spd_avg.append(SPEED[iz][m].mean())
        area.append(m.sum() * CELL_A)
    ax = axes[0]
    ax.plot(uz_avg, zs, "-o", ms=3, color="#39d0d8", label="面平均 <uz>")
    ax.plot(spd_avg, zs, "-s", ms=3, color="#ff7b72", label="面平均 |U|")
    ax.axvline(0, color=MUTED, lw=0.8)
    ax.axhline(0.0, color="#30363d", lw=0.8)
    ax.set_xlabel("速度 (m/s)", fontsize=10); ax.set_ylabel("高度 z (m)", fontsize=10)
    ax.set_title("(a) 沿高度气速衰减", fontsize=12); ax.legend(fontsize=9); ax.grid(alpha=0.2)

    ax = axes[1]
    ax.plot(qz, zs, "-o", ms=3, color=ACCENT, label="通过水平面的净风量")
    ax.axvline(10000, color="#ff7b72", ls="--", lw=1.4, label="总风量 10000 m³/h")
    ax.set_xlabel("风量 (m³/h)", fontsize=10); ax.set_ylabel("高度 z (m)", fontsize=10)
    ax.set_title("(b) 沿高度风量守恒（连续性）", fontsize=12)
    ax.legend(fontsize=9); ax.grid(alpha=0.2)
    fig.suptitle("沿高度气速衰减与风量守恒", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.96])
    fig.savefig(os.path.join(OUT, "fbed-decay.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-decay.png")


def fig_zoom_sections():
    fig, axes = plt.subplots(4, 1, figsize=(11.0, 16.0))
    for k, ax in enumerate(axes):
        quiver_slice(ax, "xz", yi(0.0), step=2, title=f"第 {k+1} 段  x∈[{0.6*k:.1f}, {0.6*(k+1):.1f}] m")
        ax.set_xlim(0.6 * k - 0.02, 0.6 * (k + 1) + 0.02)
    fig.suptitle("按区段剖分的 x-z 局部放大（含风箱进风/布风板/沸腾段）", fontsize=15, y=0.995)
    fig.tight_layout(rect=[0, 0, 1, 0.975])
    fig.savefig(os.path.join(OUT, "fbed-zoom-sections.png"), dpi=500)
    plt.close(fig)
    print("  -> fbed-zoom-sections.png")


def fig_pressure_drop():
    fig = plt.figure(figsize=(16.0, 10.5))
    gs = fig.add_gridspec(2, 2)
    ax = fig.add_subplot(gs[0, 0])
    ext = [X0, X1, Z0, Z1]
    Pxz = np.where(SOLID[:, yi(0), :], np.nan, P[:, yi(0), :])
    lim = np.nanpercentile(np.abs(Pxz), 99)
    im = ax.imshow(Pxz, origin="lower", extent=ext, cmap="RdBu_r", vmin=-lim, vmax=lim,
                   aspect="equal", interpolation="bilinear")
    ax.axhline(0, color="#7ee787", lw=1.0, ls="--")
    ax.set_title("(a) 静压场 x-z (y=0)，绿线=布风板", fontsize=12)
    ax.set_xlabel("x (m)"); ax.set_ylabel("z (m)")
    cb = fig.colorbar(im, ax=ax, fraction=0.036, pad=0.02); cb.set_label("Pa", fontsize=9)

    ax = fig.add_subplot(gs[0, 1])
    ext = [X0, X1, Y0, Y1]
    Pxy = np.where(SOLID[zi(0.10)], np.nan, P[zi(0.10)])
    im = ax.imshow(Pxy, origin="lower", extent=ext, cmap="RdBu_r",
                   vmin=-np.nanpercentile(np.abs(Pxy), 99), vmax=np.nanpercentile(np.abs(Pxy), 99),
                   aspect="equal", interpolation="bilinear")
    ax.set_title("(b) 布风板正上方 z=0.10 静压", fontsize=12)
    ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)")
    cb = fig.colorbar(im, ax=ax, fraction=0.036, pad=0.02); cb.set_label("Pa", fontsize=9)

    ax = fig.add_subplot(gs[1, 0])
    dps = []
    for k in range(N_SEC):
        ix0, ix1 = xi(0.6 * k), xi(0.6 * (k + 1))
        mp = FLUID[zi(-0.30):zi(-0.05), :, ix0:ix1]
        pp = P[zi(-0.30):zi(-0.05), :, ix0:ix1][mp].mean()
        mf = FLUID[zi(0.05):zi(0.20), :, ix0:ix1]
        pf = P[zi(0.05):zi(0.20), :, ix0:ix1][mf].mean()
        dps.append(pp - pf)
    bars = ax.bar([f"第{k+1}段" for k in range(N_SEC)], dps, color="#ff7b72", alpha=0.85)
    ax.axhline(2900, color="#7ee787", ls="--", lw=1.4, label="设计压差 2900 Pa")
    for b, v in zip(bars, dps):
        ax.text(b.get_x() + b.get_width() / 2, v + 5, f"{v:.0f}", ha="center", fontsize=9)
    ax.set_ylabel("布风板压差 (Pa)", fontsize=10); ax.legend(fontsize=9)
    ax.set_title("(c) 各段布风板压差（模型值）", fontsize=12); ax.grid(alpha=0.2, axis="y")

    ax = fig.add_subplot(gs[1, 1])
    ext = [X0, X1, Z0, Z1]
    Pxz2 = np.where(SOLID[:, yi(0), :], np.nan, P[:, yi(0), :])
    ax.plot(XC, np.nanmean(P[:, yi(0), :], axis=0), color="#39d0d8", lw=1.4, label="沿 x 平均 (y=0)")
    for zc, col in [(zi(-0.30), "#58a6ff"), (zi(-0.05), "#7ee787"), (zi(0.10), "#ff7b72")]:
        ax.plot(XC, np.where(SOLID[zc, yi(0), :], np.nan, P[zc, yi(0), :]), lw=1.0,
                color=col, label=f"z={ZC[zc]:.2f} m")
    ax.set_xlabel("x (m)", fontsize=10); ax.set_ylabel("静压 (Pa)", fontsize=10)
    ax.set_title("(d) 静压沿长度剖面", fontsize=12); ax.legend(fontsize=8); ax.grid(alpha=0.2)
    fig.suptitle("压差分析：布风板上下压差与静压分布", fontsize=15, y=0.99)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(os.path.join(OUT, "fbed-pressure-drop.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-pressure-drop.png")


def fig_flowrate():
    fig, axes = plt.subplots(1, 3, figsize=(17.0, 5.4))
    qs = []
    for k in range(N_SEC):
        ix0, ix1 = xi(0.6 * k), xi(0.6 * (k + 1))
        q = UZ[IZP, :, ix0:ix1][FLUID[IZP, :, ix0:ix1]].sum() * CELL_A * 3600.0
        qs.append(q)
    ax = axes[0]
    ax.bar([f"第{k+1}段" for k in range(N_SEC)], qs, color="#39d0d8", alpha=0.85)
    ax.axhline(2500, color="#ff7b72", ls="--", lw=1.4, label="目标 2500")
    ax.set_ylabel("风量 (m³/h)", fontsize=10); ax.legend(fontsize=9)
    ax.set_title("(a) 各区段布风板风量", fontsize=12); ax.grid(alpha=0.2, axis="y")

    ax = axes[1]
    cum = np.cumsum([UZ[IZP, :, i][FLUID[IZP, :, i]].sum() for i in range(NX)]) * CELL_A * 3600.0
    ax.plot(XC, cum, color=ACCENT, lw=1.8)
    ax.axhline(10000, color="#ff7b72", ls="--", lw=1.2, label="总计 10000")
    ax.set_xlabel("x (m)", fontsize=10); ax.set_ylabel("累计风量 (m³/h)", fontsize=10)
    ax.set_title("(b) 沿长度累计风量", fontsize=12); ax.legend(fontsize=9); ax.grid(alpha=0.2)

    ax = axes[2]
    iz_in = yi(fs.PIPE_Y0)
    qin = UY[:, iz_in, :][FLUID[:, iz_in, :]].sum() * CELL_A * 3600.0
    iz_out = NZ - 2
    qout = UZ[iz_out][FLUID[iz_out]].sum() * CELL_A * 3600.0
    qplate = sum(qs)
    labels = ["进风\n(DN300×4)", "布风板\n(通过)", "排风\n(DN800)"]
    vals = [qin, qplate, qout]
    cols = ["#7ee787", "#39d0d8", "#ff7b72"]
    b = ax.bar(labels, vals, color=cols, alpha=0.85)
    for bb, v in zip(b, vals):
        ax.text(bb.get_x() + bb.get_width() / 2, v + 60, f"{v:.0f}", ha="center", fontsize=9.5)
    ax.axhline(10000, color=MUTED, ls=":", lw=1)
    ax.set_ylabel("风量 (m³/h)", fontsize=10)
    ax.set_title("(c) 进风 / 布风板 / 排风 风量平衡", fontsize=12); ax.grid(alpha=0.2, axis="y")
    fig.suptitle("流量分析", fontsize=15, y=1.0)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(os.path.join(OUT, "fbed-flowrate.png"), dpi=520)
    plt.close(fig)
    print("  -> fbed-flowrate.png")


def main():
    fig_geometry()
    fig_vectors()
    fig_tilt()
    fig_uniformity()
    fig_decay()
    fig_zoom_sections()
    fig_pressure_drop()
    fig_flowrate()
    print("[analysis] done")


if __name__ == "__main__":
    main()