#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
规划图生成器: 系统架构图 / 数据流水线图 / 里程碑路线图
输出 docs/assets/*.svg (矢量源文件), 再由 tools/svg2png.js 转 PNG。

运行:
    python tools/gen_diagrams.py
"""
import os

OUT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                   "docs", "assets")
os.makedirs(OUT, exist_ok=True)

BG      = "#0d1117"
PANEL   = "#161b22"
PANEL2  = "#0d1117"
STROKE  = "#30363d"
FG      = "#e6edf3"
MUTED   = "#8b949e"
FONT    = "Microsoft YaHei, Segoe UI, sans-serif"

C_PINK   = "#f778ba"
C_PURPLE = "#bc8cff"
C_YELLOW = "#d29922"
C_BLUE   = "#58a6ff"
C_GREEN  = "#3fb950"
C_RED    = "#ff7b72"
C_CYAN   = "#39d0d8"


def esc(s):
    return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def text(x, y, s, size=14, fill=FG, weight="normal", anchor="start", op=1.0):
    return (f'<text x="{x}" y="{y}" font-family="{FONT}" font-size="{size}" '
            f'fill="{fill}" font-weight="{weight}" text-anchor="{anchor}" '
            f'opacity="{op}">{esc(s)}</text>\n')


def box(x, y, w, h, fill=PANEL, stroke=STROKE, rx=10, sw=1.2, dash=None, op=1.0):
    d = f' stroke-dasharray="{dash}"' if dash else ""
    return (f'<rect x="{x}" y="{y}" width="{w}" height="{h}" rx="{rx}" '
            f'fill="{fill}" stroke="{stroke}" stroke-width="{sw}"{d} opacity="{op}"/>\n')


def arrow(x1, y1, x2, y2, color=C_BLUE, sw=2.0, dash=None, marker="arr"):
    d = f' stroke-dasharray="{dash}"' if dash else ""
    return (f'<line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="{color}" '
            f'stroke-width="{sw}"{d} marker-end="url(#{marker})"/>\n')


def svg_open(w, h, extra_defs=""):
    return (f'<svg xmlns="http://www.w3.org/2000/svg" width="{w}" height="{h}" '
            f'viewBox="0 0 {w} {h}">\n'
            f'<defs>\n'
            f'<marker id="arr" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" '
            f'markerHeight="7" orient="auto-start-reverse">'
            f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{C_BLUE}"/></marker>\n'
            f'<marker id="arrG" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" '
            f'markerHeight="7" orient="auto-start-reverse">'
            f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{C_GREEN}"/></marker>\n'
            f'<marker id="arrY" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" '
            f'markerHeight="7" orient="auto-start-reverse">'
            f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{C_YELLOW}"/></marker>\n'
            f'<marker id="arrM" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" '
            f'markerHeight="7" orient="auto-start-reverse">'
            f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{MUTED}"/></marker>\n'
            f'{extra_defs}</defs>\n'
            f'<rect width="{w}" height="{h}" fill="{BG}"/>\n')


# ==========================================================================
# 1. 系统架构图
# ==========================================================================
def architecture():
    W, H = 1600, 1150
    s = svg_open(W, H)
    s += text(50, 56, "PhysicsAiFluentAnalysis · 系统架构", 28, FG, "bold")
    s += text(50, 88, "六层架构：从高保真真值到实时交互应用，物理约束贯穿模型层", 14.5, MUTED)

    layers = [
        (C_PINK,   "L5 应用层",   "Application",
         [("数字孪生", "传感器同化 · 在线重建"), ("工程设计优化", "实时参数扫描 · 设计空间探索"),
          ("教学与科研", "流场可视化实验"), ("实时决策", "风场 / 流场秒级预测")]),
        (C_PURPLE, "L4 交互与渲染", "Interaction & Rendering",
         [("WebGPU 前端", "three.js / VTK.js"), ("参数实时调节", "拖拽 Re / 攻角 / 几何"),
          ("场可视化", "体渲染 · 流线 · 矢量场"), ("场景管理", "剖切面 · 时间轴 · 导出")]),
        (C_YELLOW, "L3 推理服务层", "Inference Service",
         [("推理引擎", "TensorRT / ONNX Runtime"), ("时间推进", "自回归 + 潜空间动力学"),
          ("后处理校正", "散度修正 · 守恒律"), ("服务接口", "gRPC + WebSocket 流式")]),
        (C_BLUE,   "L2 AI 模型层", "AI Models",
         [("FNO 傅里叶神经算子", "主力基线 · 分辨率无关"), ("DeepONet", "算子学习备选路线"),
          ("PINN 物理约束", "NS 残差 · 少数据工况"), ("MeshGraphNet", "非结构网格 · 复杂几何")]),
        (C_GREEN,  "L1 数据与训练", "Data & Training",
         [("PyTorch 分布式训练", "AMP · DDP / FSDP"), ("主动学习", "不确定性驱动采样"),
          ("数据湖", "HDF5 / VTK / Zarr"), ("数据工程", "归一化 · 增广 · 插值统一")]),
        (C_RED,    "L0 高保真真值", "Ground Truth CFD",
         [("Ansys Fluent", "RANS / URANS / LES"), ("PyFluent 自动化", "参数化建模与求解"),
          ("参数化几何", "翼型族 · 车体 · 管道"), ("边界条件库", "工况模板 · 湍流模型")]),
    ]

    y0, bh, gap = 118, 148, 20
    pitch = bh + gap
    bx, bw = 230, 317
    for i, (color, name, en, mods) in enumerate(layers):
        y = y0 + i * pitch
        s += box(40, y, 1520, bh, PANEL, STROKE, 12)
        s += f'<rect x="40" y="{y}" width="6" height="{bh}" rx="3" fill="{color}"/>\n'
        s += text(64, y + 66, name, 17.5, color, "bold")
        s += text(64, y + 92, en, 11, MUTED)
        for j, (t, sub) in enumerate(mods):
            x = bx + j * (bw + 14)
            s += box(x, y + 22, bw, 104, PANEL2, STROKE, 10)
            s += f'<rect x="{x}" y="{y+22}" width="{bw}" height="4" rx="2" fill="{color}" opacity="0.75"/>\n'
            s += text(x + 16, y + 62, t, 15, FG, "bold")
            s += text(x + 16, y + 90, sub, 11.5, MUTED)
        if i < len(layers) - 1:
            ym = y + bh + gap / 2
            s += arrow(880, ym + 7, 880, ym - 7, C_BLUE, 1.8)
            s += arrow(760, ym - 7, 760, ym + 7, C_GREEN, 1.8, marker="arrG")

    s += text(40, 1128, "数据流 ↑（真值 → 训练 → 推理 → 应用）    "
                        "控制流 ↓（应用需求 → 模型选型 → 数据采集）    "
                        "物理约束：连续性 / 动量守恒 · 不确定性量化 · 守恒校正",
              12.5, MUTED)
    s += "</svg>\n"
    open(os.path.join(OUT, "architecture.svg"), "w", encoding="utf-8").write(s)
    print("  -> architecture.svg")


# ==========================================================================
# 2. 数据与训练流水线
# ==========================================================================
def pipeline():
    W, H = 1600, 760
    s = svg_open(W, H)
    s += text(50, 56, "数据与训练流水线", 28, FG, "bold")
    s += text(50, 88, "从参数采样到实时仿真服务：全链路自动化 + 主动学习数据飞轮", 14.5, MUTED)

    stages = [
        (C_CYAN,   "① 参数空间采样", ["拉丁超立方采样", "Re / 攻角 / 几何参数"]),
        (C_RED,    "② Fluent 批量仿真", ["PyFluent 无人值守", "稳态 + 瞬态工况"]),
        (C_GREEN,  "③ 数据工程", ["网格统一 · 归一化", "质检 / 增广 / 划分"]),
        (C_BLUE,   "④ 神经算子训练", ["FNO / DeepONet / GNN", "物理约束损失"]),
        (C_PURPLE, "⑤ 推理引擎部署", ["ONNX → TensorRT", "FP16 / INT8 量化"]),
        (C_PINK,   "⑥ 实时仿真服务", ["gRPC + WebSocket", "流场流式推送"]),
    ]
    arts = ["工况表 + 元数据", "CAS / DAT / VTK", "HDF5 数据集", "模型权重 .pt", "TensorRT .engine", "实时流场帧"]

    bw, gap = 224, 24
    x0, y = 68, 150
    bh = 150
    for i, (color, title, lines) in enumerate(stages):
        x = x0 + i * (bw + gap)
        s += box(x, y, bw, bh, PANEL, color, 12, sw=1.6)
        s += f'<rect x="{x}" y="{y}" width="{bw}" height="5" rx="2.5" fill="{color}"/>\n'
        s += text(x + 16, y + 44, title, 16, color, "bold")
        for k, ln in enumerate(lines):
            s += text(x + 16, y + 76 + k * 24, ln, 12, MUTED)
        if i < len(stages) - 1:
            s += arrow(x + bw + 2, y + bh / 2, x + bw + gap - 3, y + bh / 2, C_BLUE, 2.0)
        s += box(x + 8, y + bh + 16, bw - 16, 42, PANEL2, STROKE, 8, dash="4 3")
        s += text(x + bw / 2, y + bh + 43, arts[i], 11.5, MUTED, anchor="middle")

    ly = y + bh + 96
    s += arrow(x0 + 5 * (bw + gap) + bw / 2, ly - 12, x0 + bw / 2, ly - 12,
               C_YELLOW, 1.8, dash="7 5", marker="arrY")
    s += text(W / 2, ly + 16, "主动学习闭环：不确定性评估 → 高价值工况回灌仿真 → 再训练（数据飞轮）",
              13, C_YELLOW, anchor="middle")

    gy = 600
    gates = [("精度门禁", "相对 L2 < 5%"), ("稳定性门禁", "1000 步不发散"),
             ("成本门禁", "单工况预算达标"), ("可复现门禁", "流水线一键复现")]
    gw = 340
    for i, (t, sub) in enumerate(gates):
        x = 68 + i * (gw + 20)
        s += box(x, gy, gw, 66, PANEL, STROKE, 10)
        s += f'<circle cx="{x+26}" cy="{gy+33}" r="7" fill="{C_GREEN}" opacity="0.85"/>\n'
        s += text(x + 46, gy + 30, t, 14.5, FG, "bold")
        s += text(x + 46, gy + 52, sub, 11.5, MUTED)
    s += text(50, 726, "每个阶段设置 Go / No-Go 决策点：未通过门禁不允许进入下一阶段",
              12.5, MUTED)
    s += "</svg>\n"
    open(os.path.join(OUT, "pipeline.svg"), "w", encoding="utf-8").write(s)
    print("  -> pipeline.svg")


# ==========================================================================
# 3. 里程碑路线图
# ==========================================================================
def roadmap():
    W, H = 1600, 940
    s = svg_open(W, H)
    s += text(50, 56, "里程碑路线图（12 个月 · 7 个里程碑）", 28, FG, "bold")
    s += text(50, 88, "先窄后宽：2D 基准打通全链路 → 3D 复杂几何 → 实时系统 → 数字孪生", 14.5, MUTED)

    x0, x1 = 300, 1560
    mw = (x1 - x0) / 12
    ty = 140
    for m in range(13):
        x = x0 + m * mw
        s += f'<line x1="{x}" y1="{ty-24}" x2="{x}" y2="840" stroke="#21262d" stroke-width="1"/>\n'
    for m in range(12):
        s += text(x0 + (m + 0.5) * mw, ty - 6, f"{m+1}", 12.5, MUTED, anchor="middle")
    s += text(x0 - 16, ty - 6, "月份", 12.5, MUTED, anchor="end")

    miles = [
        (C_RED,    "M0 立项与调研", 1, 1, "规划文档 · 选型评审 · 基准定义"),
        (C_CYAN,   "M1 数据工厂", 2, 3, "PyFluent 自动化 · ≥500 工况数据集"),
        (C_BLUE,   "M2 基线模型", 4, 6, "FNO 基线 · 相对 L2 < 5% · 1000 步稳定"),
        (C_PURPLE, "M3 数据飞轮", 6, 7, "主动学习 · 同等精度省 40% 仿真量"),
        (C_GREEN,  "M4 3D 与复杂几何", 7, 9, "翼型 / 车体代理 · MeshGraphNet"),
        (C_YELLOW, "M5 实时推理系统", 9, 10, "TensorRT · 端到端 < 100 ms"),
        (C_PINK,   "M6 可视化与数字孪生", 10, 12, "WebGPU 交互 · 稀疏观测同化"),
    ]
    ry, rh, pitch = 170, 58, 96
    for i, (color, name, m0, m1, deliv) in enumerate(miles):
        y = ry + i * pitch
        s += text(40, y + 30, name, 16, FG, "bold")
        bx = x0 + (m0 - 1) * mw + 3
        bwid = (m1 - m0 + 1) * mw - 6
        s += box(bx, y, bwid, rh, color, color, 10, sw=0, op=0.28)
        s += f'<rect x="{bx}" y="{y}" width="{bwid}" height="{rh}" rx="10" fill="none" stroke="{color}" stroke-width="1.6"/>\n'
        span = f"第 {m0} 月" if m0 == m1 else f"{m0}-{m1} 月"
        s += text(bx + 16, y + 35, span, 13, color, "bold")
        est_w = sum(12.5 if ord(c) > 127 else 7 for c in deliv)
        if bx + bwid + 40 + est_w > 1580:      # 右侧放不下 -> 放左边
            dx = bx - 14
            s += text(dx - 16, y + rh / 2 + 5, deliv, 12.5, MUTED, anchor="end")
            s += f'<path d="M {dx} {y+rh/2-7} L {dx-7} {y+rh/2} L {dx} {y+rh/2+7} L {dx+7} {y+rh/2} z" fill="{color}"/>\n'
        else:
            dx = bx + bwid + 14
            s += f'<path d="M {dx} {y+rh/2-7} L {dx+7} {y+rh/2} L {dx} {y+rh/2+7} L {dx-7} {y+rh/2} z" fill="{color}"/>\n'
            s += text(dx + 18, y + rh / 2 + 5, deliv, 12.5, MUTED)

    s += box(40, 856, 1520, 60, PANEL, STROKE, 10)
    s += text(60, 884, "决策点", 14, C_YELLOW, "bold")
    s += text(130, 884, "M2 精度不达标 → DeepONet / MeshGraphNet 双路线并行    "
                        "M4 3D 不达标 → 退守「2D + 切片 3D」保证 M5/M6 交付    "
                        "每里程碑预留 20% 时间缓冲", 12.5, MUTED)
    s += text(60, 906, "成功标准", 14, C_GREEN, "bold")
    s += text(130, 906, "P1（3 个月）可行性验证 · P2（6 个月）系统闭环 · P3（12 个月）数字孪生原型",
              12.5, MUTED)
    s += "</svg>\n"
    open(os.path.join(OUT, "roadmap.svg"), "w", encoding="utf-8").write(s)
    print("  -> roadmap.svg")


if __name__ == "__main__":
    architecture()
    pipeline()
    roadmap()
    print("done:", OUT)