#!/usr/bin/env python3
"""
visualize_sdf.py
────────────────────────────────────────────────────────────────────────
读取 dump_sdf_slice() 生成的文本文件，绘制：
  · 各 subface 的 SDF 热力图（蓝=负/内部，红=正/外部，白=零面）
  · 零等值线（石灰绿，代表算法"认定"的模型表面）
  · 合成 SDF（max 合并，即 CSG 求交结果）

用法：
  pip install numpy matplotlib scipy
  python visualize_sdf.py sdf_slice.txt [y_slice_for_xy_view]
────────────────────────────────────────────────────────────────────────
"""

import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker
from pathlib import Path
from scipy.ndimage import gaussian_filter
from matplotlib.gridspec import GridSpec

# ── surface_type 枚举（与 C++ 侧对齐）────────────────────────────────
SURFACE_NAMES = {
    0: "Plane (cap)",
    1: "Sphere",
    2: "Cylinder",
    3: "Cone",
    4: "ExtrudePolyline\nSide Face",
    5: "ExtrudeHelixline\nSide Face",
}

# 设置全局样式
plt.rcParams.update({
    'font.size': 9,
    'axes.titlesize': 10,
    'axes.labelsize': 9,
    'xtick.labelsize': 8,
    'ytick.labelsize': 8,
    'legend.fontsize': 8,
    'figure.titlesize': 12,
    'figure.autolayout': False,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.1,
})

# ─────────────────────────────────────────────────────────────────────
# 1. 读取数据
# ─────────────────────────────────────────────────────────────────────
def load_sdf_data(path: str) -> dict:
    """解析 dump_sdf_slice 输出文件，返回结构化字典。"""
    lines = Path(path).read_text().splitlines()
    idx = 0

    # 行 0: grid_res y_slice
    tok = lines[idx].split(); idx += 1
    grid_res = int(tok[0])
    y_slice  = float(tok[1])

    # 行 1: x0 x1 z0 z1
    tok = lines[idx].split(); idx += 1
    x0, x1, z0, z1 = float(tok[0]), float(tok[1]), float(tok[2]), float(tok[3])

    # 行 2: num_subfaces
    ns = int(lines[idx]); idx += 1

    N = grid_res + 1
    subfaces = []
    for _ in range(ns):
        stype = int(lines[idx]); idx += 1
        rows = []
        for _ in range(N):
            rows.append(list(map(float, lines[idx].split())))
            idx += 1
        subfaces.append({"type": stype, "grid": np.array(rows, dtype=np.float64)})

    return dict(grid_res=grid_res, y_slice=y_slice,
                x0=x0, x1=x1, z0=z0, z1=z1, subfaces=subfaces)


# ─────────────────────────────────────────────────────────────────────
# 2. 绘制单个面板
# ─────────────────────────────────────────────────────────────────────
def draw_panel(ax, XX, ZZ, grid: np.ndarray, title: str, y_slice: float,
               cmap=None, show_neg_region=True, highlight_sign_flip=False):
    """
    在 ax 上绘制：
      · 连续色带热力图（双斜率归一化，保证零值白色）
      · 石灰绿零等值线（模型表面位置）
      · 负值区边界（蓝色虚线，"算法认为的内部"边界）
    """
    if cmap is None:
        cmap = plt.cm.RdBu_r  # 改为 RdBu_r，红色为正，蓝色为负，更直观

    finite = grid[np.isfinite(grid)]
    if finite.size == 0:
        ax.set_title(title + "\n[no finite data]")
        return

    # 以 98 分位数作为色域上限，避免极值压缩色彩
    vmax = float(np.percentile(np.abs(finite), 98))
    vmax = max(vmax, 1e-6)
    norm = mcolors.TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax)

    # 应用轻微高斯平滑，使热力图更清晰
    if grid.shape[0] > 10 and grid.shape[1] > 10:
        grid_smoothed = gaussian_filter(grid, sigma=0.7, mode='nearest')
    else:
        grid_smoothed = grid

    # ── 填充等高线热力图 ──
    cf = ax.contourf(XX, ZZ, grid_smoothed, levels=50, cmap=cmap, norm=norm, 
                     alpha=0.85, extend='both')
    
    # 添加颜色条，但位置更紧凑
    cbar = plt.colorbar(cf, ax=ax, pad=0.01, fraction=0.045, shrink=0.9)
    cbar.set_label("SDF", fontsize=7, labelpad=2)
    cbar.ax.tick_params(labelsize=6, pad=1)

    # ── 零等值线（模型表面） ──
    try:
        cs0 = ax.contour(XX, ZZ, grid_smoothed, levels=[0.0],
                         colors=["lime"], linewidths=1.8, zorder=5, alpha=0.9)
        if len(cs0.collections) > 0:
            # 只在有零等值线的地方添加标签
            ax.clabel(cs0, fmt=" SDF=0 ", fontsize=7, inline=True,
                      colors="lime", zorder=6, inline_spacing=5)
    except Exception:
        pass

    # 高亮显示符号突变区域（如果启用）
    if highlight_sign_flip and finite.size > 4:
        from scipy.ndimage import generic_filter
        def local_range(vals):
            return vals.max() - vals.min() if vals.size > 0 else 0
        
        # 检测局部范围内的剧烈符号变化
        local_range_grid = generic_filter(grid, local_range, size=3)
        mean_range = float(np.percentile(local_range_grid, 85))
        sign_flip_mask = (local_range_grid > mean_range * 2.5) & np.isfinite(grid)
        
        if sign_flip_mask.any():
            # 用半透明黄色填充显示符号突变区域
            ax.contourf(XX, ZZ, sign_flip_mask.astype(float), 
                       levels=[0.5, 1.5], colors=["yellow"], 
                       alpha=0.25, zorder=4)
            ax.contour(XX, ZZ, sign_flip_mask.astype(float), levels=[0.5],
                      colors=["yellow"], linewidths=1.0, linestyles="dotted", 
                      zorder=5, alpha=0.7)
            
            # 添加图例说明
            sign_flip_pct = sign_flip_mask.sum() / sign_flip_mask.size * 100
            if sign_flip_pct > 0.5:  # 仅在显著区域添加标注
                ax.text(0.02, 0.98, f"⚠ {sign_flip_pct:.1f}% sign-flip", 
                        transform=ax.transAxes, fontsize=6, color="black",
                        bbox=dict(boxstyle="round,pad=0.2", fc="yellow", 
                                 alpha=0.7, ec="black", lw=0.5),
                        zorder=10, verticalalignment='top')

    # ── 统计信息标注 ──
    neg_pct = np.sum(grid < 0) / grid.size * 100
    pos_pct = 100 - neg_pct
    min_val, max_val = np.min(grid), np.max(grid)
    
    # 更简洁的统计信息
    info = f"Neg: {neg_pct:.1f}%\nPos: {pos_pct:.1f}%\nMin: {min_val:.2e}\nMax: {max_val:.2e}"
    ax.text(0.02, 0.02, info, transform=ax.transAxes,
            fontsize=6, color="white", zorder=7,
            bbox=dict(boxstyle="round,pad=0.2", fc="black", alpha=0.6, 
                     ec="gray", lw=0.5),
            verticalalignment='bottom')

    # ── 坐标轴 ──
    ax.set_xlabel("X", fontsize=8, labelpad=2)
    ax.set_ylabel("Z", fontsize=8, labelpad=2)
    ax.set_title(title, fontsize=9, pad=6, fontweight='medium')
    ax.set_aspect("equal", adjustable="box")
    ax.tick_params(labelsize=7, pad=2)
    ax.xaxis.set_minor_locator(mticker.AutoMinorLocator(2))
    ax.yaxis.set_minor_locator(mticker.AutoMinorLocator(2))
    
    # 更精细的网格
    ax.grid(True, which="major", color="gray", alpha=0.2, lw=0.3, linestyle='-')
    ax.grid(True, which="minor", color="gray", alpha=0.1, lw=0.1, linestyle=':')


# ─────────────────────────────────────────────────────────────────────
# 3. 主绘图入口 - 重新设计布局
# ─────────────────────────────────────────────────────────────────────
def plot_sdf_heatmap(data: dict, out_path: str):
    subfaces = data["subfaces"]
    ns = len(subfaces)
    y_slice = data["y_slice"]
    x0, x1, z0, z1 = data["x0"], data["x1"], data["z0"], data["z1"]
    N = data["grid_res"] + 1

    X = np.linspace(x0, x1, N)
    Z = np.linspace(z0, z1, N)
    XX, ZZ = np.meshgrid(X, Z)

    # 计算合成SDF
    if ns > 0:
        combined = subfaces[0]["grid"].copy()
        for sf in subfaces[1:]:
            combined = np.maximum(combined, sf["grid"])
    
    # 智能布局计算
    if ns <= 1:
        # 只有1个子面：一行两列
        fig = plt.figure(figsize=(12, 5.5))
        gs = GridSpec(1, 3, width_ratios=[1, 0.05, 1.2], wspace=0.15, hspace=0.1)
        axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 2])]
    elif ns == 2:
        # 2个子面：一行三列
        fig = plt.figure(figsize=(15, 5))
        gs = GridSpec(1, 4, width_ratios=[1, 1, 0.05, 1.2], wspace=0.15, hspace=0.1)
        axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 3])]
    elif ns == 3:
        # 3个子面：两行两列
        fig = plt.figure(figsize=(12, 9))
        gs = GridSpec(2, 3, width_ratios=[1, 1, 0.05], height_ratios=[1, 1], wspace=0.15, hspace=0.2)
        axes = [
            fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]),
            fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]),
            fig.add_subplot(gs[0:2, 2])  # 合成图占两行
        ]
    else:
        # 更多子面：计算最优布局
        ncols = min(3, int(np.ceil(np.sqrt(ns + 1))))
        nrows = int(np.ceil((ns + 1) / ncols))
        fig_width = min(6 * ncols, 18)
        fig_height = 5 * nrows
        
        fig, axes = plt.subplots(nrows, ncols, 
                                figsize=(fig_width, fig_height),
                                squeeze=False)
        axes = axes.flatten()

    cmap = plt.cm.RdBu_r

    # 绘制每个子面
    for s, sf in enumerate(subfaces):
        if s < len(axes) - 1:  # 最后一个位置留给合成图
            stype = sf["type"]
            name = SURFACE_NAMES.get(stype, f"Type {stype}")
            title = f"Subface {s}: {name}"
            
            # 对于第4种类型（ExtrudePolyline Side Face），启用符号突变检测
            highlight = (stype == 4)
            
            draw_panel(axes[s], XX, ZZ, sf["grid"], title, y_slice, 
                       cmap, highlight_sign_flip=highlight)

    # 绘制合成面板
    if ns > 0 and ns < len(axes):
        title_comb = "Combined SDF\nmax(subfaces) ≈ CSG Intersection"
        draw_panel(axes[ns], XX, ZZ, combined, title_comb, y_slice, cmap, 
                   highlight_sign_flip=True)
        
        # 标记隐藏多余子图
        for i in range(ns + 1, len(axes)):
            axes[i].set_visible(False)
    elif ns == 0:
        axes[0].text(0.5, 0.5, "No subface data", 
                    ha='center', va='center', transform=axes[0].transAxes)
        axes[0].set_title("Empty Data")
        for i in range(1, len(axes)):
            axes[i].set_visible(False)

    # 添加全局标题
    fig.suptitle(
        f"SDF Cross-Section Analysis - Y = {y_slice:.6f}\n"
        "Blue = Inside Model (SDF < 0) | Red = Outside Model (SDF > 0) | "
        "Lime Green = Model Surface (SDF = 0)",
        fontsize=11, y=0.98, fontweight='bold'
    )

    # 添加图例说明
    fig.text(0.02, 0.02, 
             f"Data: {Path(out_path).stem}.txt | Grid: {data['grid_res']}×{data['grid_res']} | "
             f"Subfaces: {ns} | X: [{x0:.3f}, {x1:.3f}] | Z: [{z0:.3f}, {z1:.3f}]",
             fontsize=7, style='italic', alpha=0.7)

    # 保存图形
    fig.savefig(out_path, dpi=200, bbox_inches="tight", facecolor='white')
    print(f"[SDF_VIZ] Heatmap saved → {out_path}")
    
    # 显示图形
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])  # 为标题留出空间
    plt.show()


# ─────────────────────────────────────────────────────────────────────
# 4. CLI 入口
# ─────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python visualize_sdf.py <sdf_data_file.txt>")
        print("Example: python visualize_sdf.py sdf_slice.txt")
        sys.exit(1)
    
    data_path = sys.argv[1]
    
    if not Path(data_path).exists():
        print(f"[ERROR] File not found: {data_path}")
        sys.exit(1)

    print(f"[SDF_VIZ] Loading {data_path} ...")
    try:
        data = load_sdf_data(data_path)
        print(f"[SDF_VIZ] Grid: {data['grid_res']}×{data['grid_res']} | "
              f"Y-slice: {data['y_slice']:.6f} | "
              f"Subfaces: {len(data['subfaces'])}")
        
        # 生成输出路径
        stem = Path(data_path).stem
        out_dir = Path(data_path).parent
        out_path = str(out_dir / f"{stem}_heatmap.png")
        
        # 绘制热力图
        plot_sdf_heatmap(data, out_path)
        
        # 打印文件信息
        if Path(out_path).exists():
            file_size = Path(out_path).stat().st_size
            print(f"[SDF_VIZ] Output: {out_path} ({file_size/1024:.1f} KB)")
        
    except Exception as e:
        print(f"[ERROR] Failed to process {data_path}: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)