#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建多Agent数据分析系统流程图
展示整个系统的运作流程和Agent间的协作关系
"""

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, Circle, Arrow
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def create_agent_flowchart():
    """创建Agent运作流程图"""
    
    # 创建图形和坐标轴
    fig, ax = plt.subplots(1, 1, figsize=(20, 16))
    ax.set_xlim(0, 20)
    ax.set_ylim(0, 16)
    ax.axis('off')
    
    # 定义颜色方案
    colors = {
        'user': '#FF6B6B',           # 用户 - 红色
        'supervisor': '#4ECDC4',      # 监督者 - 青色
        'rag': '#45B7D1',            # RAG Agent - 蓝色
        'python': '#96CEB4',         # Python Agent - 绿色
        'amap': '#FECA57',           # AMAP Agent - 黄色
        'file': '#FF9FF3',           # File Agent - 粉色
        'sql': '#54A0FF',            # SQL Agent - 深蓝色
        'terminal': '#5F27CD',       # Terminal Agent - 紫色
        'web_search': '#00D2D3',     # Web Search Agent - 青绿色
        'tools': '#DDA0DD',          # 工具 - 淡紫色
        'external': '#FFB8B8'        # 外部服务 - 浅红色
    }
    
    # 1. 用户输入部分
    user_box = FancyBboxPatch((1, 14), 3, 1.5, 
                             boxstyle="round,pad=0.1",
                             facecolor=colors['user'], 
                             edgecolor='black', linewidth=2)
    ax.add_patch(user_box)
    ax.text(2.5, 14.75, '用户查询\nUser Query', ha='center', va='center', 
            fontsize=12, fontweight='bold', color='white')
    
    # 2. Supervisor中央调度
    supervisor_box = FancyBboxPatch((8, 12), 4, 2, 
                                   boxstyle="round,pad=0.15",
                                   facecolor=colors['supervisor'], 
                                   edgecolor='black', linewidth=3)
    ax.add_patch(supervisor_box)
    ax.text(10, 13, '智能Supervisor\n任务分析与路由\n(LangGraph)', ha='center', va='center', 
            fontsize=14, fontweight='bold', color='white')
    
    # 3. 七个专业Agent
    agents = [
        {'name': 'Enhanced\nRAG Agent\n电信数据专家', 'pos': (2, 9), 'color': colors['rag']},
        {'name': 'Enhanced\nPython Agent\n数据科学专家', 'pos': (6, 9), 'color': colors['python']},
        {'name': 'Enhanced\nAMAP Agent\n地理服务专家', 'pos': (10, 9), 'color': colors['amap']},
        {'name': 'Safe File\nAgent\n文件管理专家', 'pos': (14, 9), 'color': colors['file']},
        {'name': 'SQL Agent\n数据库专家', 'pos': (18, 9), 'color': colors['sql']},
        {'name': 'Terminal\nCommand Agent\n系统操作专家', 'pos': (4, 6), 'color': colors['terminal']},
        {'name': 'Web Search\nAgent\n网络检索专家', 'pos': (12, 6), 'color': colors['web_search']}
    ]
    
    # 绘制Agent boxes
    agent_boxes = {}
    for agent in agents:
        box = FancyBboxPatch((agent['pos'][0]-1, agent['pos'][1]-0.75), 2, 1.5,
                            boxstyle="round,pad=0.1",
                            facecolor=agent['color'],
                            edgecolor='black', linewidth=2)
        ax.add_patch(box)
        ax.text(agent['pos'][0], agent['pos'][1], agent['name'], 
                ha='center', va='center', fontsize=10, fontweight='bold', color='white')
        agent_boxes[agent['name'].split('\n')[1]] = agent['pos']
    
    # 4. 外部服务和工具
    external_services = [
        {'name': 'Tavily API\n网络搜索', 'pos': (16, 4), 'color': colors['external']},
        {'name': 'AMAP API\n地理服务', 'pos': (8, 4), 'color': colors['external']},
        {'name': 'MySQL\n数据库', 'pos': (18, 4), 'color': colors['external']},
        {'name': 'Local Vector DB\n知识库', 'pos': (2, 4), 'color': colors['external']},
        {'name': 'File System\n文件系统', 'pos': (14, 4), 'color': colors['external']},
        {'name': 'Terminal\n系统命令', 'pos': (4, 4), 'color': colors['external']}
    ]
    
    for service in external_services:
        box = FancyBboxPatch((service['pos'][0]-1, service['pos'][1]-0.5), 2, 1,
                            boxstyle="round,pad=0.05",
                            facecolor=service['color'],
                            edgecolor='gray', linewidth=1)
        ax.add_patch(box)
        ax.text(service['pos'][0], service['pos'][1], service['name'], 
                ha='center', va='center', fontsize=9, color='black')
    
    # 5. 工具层
    tools = [
        {'name': 'Python\n执行环境', 'pos': (6, 2), 'color': colors['tools']},
        {'name': 'matplotlib\nseaborn', 'pos': (10, 2), 'color': colors['tools']},
        {'name': 'scikit-learn\nML模型', 'pos': (14, 2), 'color': colors['tools']}
    ]
    
    for tool in tools:
        box = FancyBboxPatch((tool['pos'][0]-0.8, tool['pos'][1]-0.4), 1.6, 0.8,
                            boxstyle="round,pad=0.05",
                            facecolor=tool['color'],
                            edgecolor='gray', linewidth=1)
        ax.add_patch(box)
        ax.text(tool['pos'][0], tool['pos'][1], tool['name'], 
                ha='center', va='center', fontsize=8)
    
    # 6. 绘制箭头连接
    # 用户到Supervisor
    ax.annotate('', xy=(8, 13), xytext=(4, 14.5),
                arrowprops=dict(arrowstyle='->', lw=3, color='red'))
    
    # Supervisor到各Agent
    supervisor_pos = (10, 12)
    for agent in agents:
        ax.annotate('', xy=agent['pos'], xytext=supervisor_pos,
                    arrowprops=dict(arrowstyle='->', lw=2, color='blue', alpha=0.7))
    
    # Agent到外部服务的连接
    connections = [
        ((2, 9), (2, 4.5)),      # RAG Agent -> Vector DB
        ((6, 9), (6, 2.5)),      # Python Agent -> Python环境
        ((10, 9), (8, 4.5)),     # AMAP Agent -> AMAP API
        ((14, 9), (14, 4.5)),    # File Agent -> File System
        ((18, 9), (18, 4.5)),    # SQL Agent -> MySQL
        ((4, 6), (4, 4.5)),      # Terminal Agent -> Terminal
        ((12, 6), (16, 4.5))     # Web Search Agent -> Tavily API
    ]
    
    for start, end in connections:
        ax.annotate('', xy=end, xytext=start,
                    arrowprops=dict(arrowstyle='->', lw=1.5, color='green', alpha=0.8))
    
    # 7. 流程步骤说明
    steps_box = FancyBboxPatch((0.5, 0.5), 7, 1,
                              boxstyle="round,pad=0.1",
                              facecolor='lightgray',
                              edgecolor='black', linewidth=1)
    ax.add_patch(steps_box)
    ax.text(4, 1, '处理流程: 1.用户输入 → 2.Supervisor分析 → 3.选择Agent → 4.执行任务 → 5.返回结果', 
            ha='center', va='center', fontsize=11, fontweight='bold')
    
    # 8. 协作说明
    collab_box = FancyBboxPatch((12.5, 0.5), 7, 1,
                               boxstyle="round,pad=0.1",
                               facecolor='lightyellow',
                               edgecolor='black', linewidth=1)
    ax.add_patch(collab_box)
    ax.text(16, 1, '协作模式: Agent间可以互相调用，形成复杂的工作流程链', 
            ha='center', va='center', fontsize=11, fontweight='bold')
    
    # 9. 添加标题
    ax.text(10, 15.5, '多Agent数据分析系统架构流程图', 
            ha='center', va='center', fontsize=18, fontweight='bold')
    
    # 10. 添加图例
    legend_elements = [
        mpatches.Rectangle((0, 0), 1, 1, facecolor=colors['user'], label='用户交互'),
        mpatches.Rectangle((0, 0), 1, 1, facecolor=colors['supervisor'], label='智能调度'),
        mpatches.Rectangle((0, 0), 1, 1, facecolor=colors['rag'], label='专业Agent'),
        mpatches.Rectangle((0, 0), 1, 1, facecolor=colors['external'], label='外部服务'),
        mpatches.Rectangle((0, 0), 1, 1, facecolor=colors['tools'], label='工具库')
    ]
    
    ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(0.98, 0.98))
    
    # 11. 添加技术栈说明
    tech_text = """
    技术栈:
    • LangGraph: 工作流编排
    • LangChain: Agent框架
    • DeepSeek: 大语言模型
    • Tavily: 网络搜索API
    • AMAP: 地理位置服务
    • MySQL: 数据存储
    • Python: 数据分析
    """
    
    ax.text(0.5, 11, tech_text, ha='left', va='top', fontsize=9, 
            bbox=dict(boxstyle="round,pad=0.5", facecolor='lightblue', alpha=0.7))
    
    plt.tight_layout()
    return fig

def main():
    """主函数"""
    print("🎨 开始创建Agent运作流程图...")
    
    # 创建流程图
    fig = create_agent_flowchart()
    
    # 保存图片
    output_path = "/Users/liaowang/Desktop/Data_Agent项目文件/agent_system_flowchart.png"
    fig.savefig(output_path, dpi=300, bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    
    print(f"✅ 流程图已保存到: {output_path}")
    print("📊 流程图包含以下内容:")
    print("  - 完整的系统架构图")
    print("  - 7个专业Agent及其功能")
    print("  - Supervisor智能调度机制")
    print("  - 外部服务和工具集成")
    print("  - 处理流程和协作模式")
    print("  - 技术栈说明")
    
    # 也保存一个高分辨率版本
    high_res_path = "/Users/liaowang/Desktop/Data_Agent项目文件/agent_system_flowchart_high_res.png"
    fig.savefig(high_res_path, dpi=600, bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    print(f"🎯 高分辨率版本已保存到: {high_res_path}")
    
    plt.close(fig)

if __name__ == "__main__":
    main()