visualization.py

"""
可视化模块 - 增强版
提供更直观、更丰富的模拟过程可视化
包含高级算法状态显示
"""
import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.lines import Line2D
from matplotlib.animation import FuncAnimation
from typing import List, Optional, Dict
import numpy as np
import platform

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
FIGURES_DIR = os.path.join(BASE_DIR, "figures")

# 配置中文字体支持
if platform.system() == 'Darwin':  # macOS
    plt.rcParams['font.family'] = ['Arial Unicode MS', 'Hiragino Sans GB', 'STHeiti', 'Songti SC', 'sans-serif']
elif platform.system() == 'Windows':
    plt.rcParams['font.family'] = ['Microsoft YaHei', 'SimHei', 'sans-serif']
else:  # Linux
    plt.rcParams['font.family'] = ['WenQuanYi Micro Hei', 'Noto Sans CJK SC', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

from src.agent_base import Position
from src.environment import Environment, SearchZone
from src.drone_agent import DroneAgent, DroneState
from src.coordinator_agent import CoordinatorAgent
from src.advanced_algorithms import AgentRole


class Visualizer:
    """
    增强可视化器
    提供更丰富的图例和更直观的显示
    """

    def __init__(self, environment: Environment):
        self.env = environment
        self.fig, (self.ax_main, self.ax_info) = plt.subplots(
            1, 2, figsize=(16, 10),
            gridspec_kw={'width_ratios': [3, 1]}
        )
        self.drone_plots = {}
        self.drone_labels = {}
        self.drone_ranges = {}
        self.target_plots = {}
        self.zone_patches = {}
        self.trail_plots = {}
        self.coverage_layer = None  # 热力图图层

        # 颜色配置 - 更鲜明的配色
        self.drone_colors = {
            0: '#E74C3C',  # 红色 - Alpha
            1: '#2ECC71',  # 绿色 - Bravo
            2: '#3498DB',  # 蓝色 - Charlie
            3: '#9B59B6',  # 紫色 - Delta
            4: '#F39C12',  # 橙色
            5: '#1ABC9C',  # 青色
        }

        self.colors = {
            'base': '#F1C40F',           # 金黄色基地
            'target_unfound': '#E74C3C',  # 红色未发现
            'target_found': '#2ECC71',    # 绿色已发现
            'obstacle': '#7F8C8D',        # 灰色障碍物
            'zone_unassigned': '#ECF0F1', # 浅灰未分配区域
            'zone_assigned': '#D5F5E3',   # 浅绿已分配区域
            'zone_completed': '#A9DFBF',  # 深绿完成区域
        }

        # 状态对应的标记
        self.state_markers = {
            DroneState.IDLE: 'o',           # 圆形
            DroneState.MOVING_TO_ZONE: '^', # 三角形
            DroneState.SEARCHING: 's',      # 方形
            DroneState.TARGET_FOUND: '*',   # 星形
            DroneState.RETURNING: 'v',      # 倒三角
            DroneState.CHARGING: 'p',       # 五边形
        }

    def setup(self, drones: List[DroneAgent]):
        """初始化可视化"""
        self._setup_main_view(drones)
        self._setup_info_panel(drones)
        plt.tight_layout()

    def _setup_main_view(self, drones: List[DroneAgent]):
        """设置主视图"""
        ax = self.ax_main
        ax.clear()
        ax.set_xlim(-5, self.env.width + 5)
        ax.set_ylim(-5, self.env.height + 5)
        ax.set_aspect('equal')
        ax.set_title('无人机集群协作搜救系统 - 实时态势', fontsize=16, fontweight='bold', pad=15)
        ax.set_xlabel('X 坐标 (米)', fontsize=11)
        ax.set_ylabel('Y 坐标 (米)', fontsize=11)
        ax.grid(True, alpha=0.3, linestyle='--')
        ax.set_facecolor('#FAFAFA')

        # 0. 绘制覆盖热力图（最底层）
        # 创建自定义colormap: 0=透明/白, 1=浅绿(已探索), 2=亮黄(热点)
        from matplotlib.colors import LinearSegmentedColormap
        colors = [(0, 0, 0, 0), (0.2, 0.8, 0.2, 0.3), (1, 0.8, 0, 0.5)]
        cmap = LinearSegmentedColormap.from_list('coverage', colors, N=3)
        
        # 初始化空的网格数据
        grid_data = np.zeros((self.env.coverage_map.rows, self.env.coverage_map.cols))
        self.coverage_layer = ax.imshow(
            grid_data, 
            extent=[0, self.env.width, 0, self.env.height], 
            origin='lower',
            cmap=cmap,
            vmin=0, vmax=2,
            zorder=0  # 最底层
        )

        # 1. 绘制搜索区域（底层）
        for i, zone in enumerate(self.env.search_zones):
            width = zone.bottom_right.x - zone.top_left.x
            height = zone.bottom_right.y - zone.top_left.y

            color = self.colors['zone_unassigned']
            if zone.searched:
                color = self.colors['zone_completed']
            elif zone.assigned_to:
                color = self.colors['zone_assigned']

            rect = patches.Rectangle(
                (zone.top_left.x, zone.top_left.y),
                width, height,
                linewidth=2, edgecolor='#333333',
                facecolor=color, alpha=0.4,
                linestyle='--'
            )
            self.zone_patches[zone.zone_id] = rect
            ax.add_patch(rect)

            # 区域标签
            ax.text(
                zone.center.x, zone.center.y,
                f'{zone.zone_id}\n优先级:{zone.priority}',
                ha='center', va='center', fontsize=9,
                alpha=0.6, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.7)
            )

        # 2. 绘制障碍物
        for obs in self.env.obstacles:
            circle = patches.Circle(
                (obs.position.x, obs.position.y), obs.radius,
                facecolor=self.colors['obstacle'],
                edgecolor='#2C3E50',
                alpha=0.6, linewidth=1.5,
                hatch='///'
            )
            ax.add_patch(circle)

        # 3. 绘制基地（在地图中心）
        base_x = self.env.base_position.x if hasattr(self.env, 'base_position') else self.env.width / 2
        base_y = self.env.base_position.y if hasattr(self.env, 'base_position') else self.env.height / 2
        base = patches.RegularPolygon(
            (base_x, base_y), numVertices=6, radius=5,
            facecolor=self.colors['base'],
            edgecolor='#D68910',
            linewidth=2, alpha=0.9
        )
        ax.add_patch(base)
        ax.text(base_x, base_y, '基地', ha='center', va='center',
                fontsize=10, fontweight='bold', color='#6E2C00')

        # 4. 绘制目标 - 根据目标类型使用不同显示
        from src.environment import TargetType
        for target in self.env.targets:
            color = self.colors['target_found'] if target.found else self.colors['target_unfound']

            # 根据目标类型调整显示
            target_type = getattr(target, 'target_type', TargetType.NORMAL)
            if target_type == TargetType.OCCLUDED:
                # 遮蔽目标：用方形标记，橙色边框
                marker = 's' if not target.found else '*'
                edge_color = '#E67E22'  # 橙色
            elif target_type == TargetType.TRAPPED:
                # 被困建筑目标：用菱形标记，红色边框
                marker = 'D' if not target.found else '*'
                edge_color = '#C0392B'  # 深红色
            elif target_type == TargetType.MOBILE:
                # 移动目标：用圆形标记，带箭头，紫色边框
                marker = 'o' if not target.found else '*'
                edge_color = '#8E44AD'  # 紫色
            else:
                # 普通目标
                marker = 'X' if not target.found else '*'
                edge_color = '#2C3E50'

            plot, = ax.plot(
                target.position.x, target.position.y,
                marker=marker,
                markersize=16 if not target.found else 18,
                color=color,
                markeredgecolor=edge_color,
                markeredgewidth=2
            )
            self.target_plots[target.target_id] = plot

            # 目标标签 - 简化显示，只显示ID和类型
            type_short = {'普通': '', '遮蔽': '[遮]', '被困建筑': '[困]'}
            type_label = type_short.get(target_type.value, '') if target_type != TargetType.NORMAL else ''
            status = '(OK)' if target.found else ''
            label_text = f'{target.target_id}{type_label}{status}'

            ax.annotate(
                label_text,
                (target.position.x, target.position.y),
                xytext=(8, -12), textcoords='offset points',
                fontsize=7, alpha=0.9,
                bbox=dict(boxstyle='round,pad=0.15', facecolor='white', alpha=0.7, edgecolor='none')
            )

        # 5. 绘制无人机
        for i, drone in enumerate(drones):
            color = self.drone_colors.get(i, '#333333')

            # 感知范围圆
            range_circle = patches.Circle(
                (drone.position.x, drone.position.y),
                drone.spec.perception_range,
                facecolor=color, alpha=0.1,
                edgecolor=color, linestyle=':', linewidth=1.5
            )
            ax.add_patch(range_circle)
            self.drone_ranges[drone.agent_id] = range_circle

            # 无人机位置
            marker = self.state_markers.get(drone.drone_state, '^')
            plot, = ax.plot(
                drone.position.x, drone.position.y,
                marker=marker, markersize=14, color=color,
                markeredgecolor='#2C3E50', markeredgewidth=2
            )
            self.drone_plots[drone.agent_id] = (plot, i)

            # 无人机标签 - 紧贴无人机
            role = getattr(drone, 'current_role', None)
            role_str = f"[{role.value}]" if role else ""
            label = ax.annotate(
                f'{drone.name}{role_str}\n{drone.drone_state.value}\n电量:{drone.battery:.0f}%',
                (drone.position.x, drone.position.y),
                xytext=(5, 5), textcoords='offset points',
                fontsize=7, fontweight='bold',
                color=color,
                bbox=dict(boxstyle='round,pad=0.2', facecolor='white',
                         edgecolor=color, alpha=0.85)
            )
            self.drone_labels[drone.agent_id] = label

            # 轨迹
            trail, = ax.plot([], [], '-', color=color, alpha=0.4, linewidth=2)
            self.trail_plots[drone.agent_id] = (trail, [])

    def _setup_info_panel(self, drones: List[DroneAgent]):
        """设置信息面板"""
        ax = self.ax_info
        ax.clear()
        ax.set_xlim(0, 10)
        ax.set_ylim(0, 24)
        ax.axis('off')
        ax.set_title('系统状态面板', fontsize=12, fontweight='bold', pad=5)

        # 创建丰富的图例 - 调整布局使其更紧凑
        y_pos = 23.5

        # === 无人机图例 ===
        ax.text(0.5, y_pos, '【无人机】', fontsize=10, fontweight='bold', color='#2C3E50')
        y_pos -= 0.8

        for i, drone in enumerate(drones):
            color = self.drone_colors.get(i, '#333333')
            ax.plot(1, y_pos, '^', markersize=8, color=color,
                   markeredgecolor='#2C3E50', markeredgewidth=1.0)
            ax.text(1.8, y_pos, drone.name, fontsize=8, va='center', color=color, fontweight='bold')
            y_pos -= 0.5

        y_pos -= 0.3

        # === 无人机状态图例 ===
        ax.text(0.5, y_pos, '【状态标记】', fontsize=10, fontweight='bold', color='#2C3E50')
        y_pos -= 0.8

        state_legend = [
            ('o', '空闲', '#95A5A6'),
            ('^', '前往区域', '#3498DB'),
            ('s', '搜索中', '#E67E22'),
            ('*', '发现目标', '#E74C3C'),
            ('v', '返航', '#9B59B6'),
            ('p', '充电中', '#2ECC71'),
        ]

        # 双列布局状态
        for i in range(0, len(state_legend), 2):
            # 左列
            marker, label, color = state_legend[i]
            ax.plot(1, y_pos, marker, markersize=7, color=color)
            ax.text(1.8, y_pos, label, fontsize=7.5, va='center')
            
            # 右列
            if i + 1 < len(state_legend):
                marker2, label2, color2 = state_legend[i+1]
                ax.plot(5.5, y_pos, marker2, markersize=7, color=color2)
                ax.text(6.3, y_pos, label2, fontsize=7.5, va='center')
            y_pos -= 0.5

        y_pos -= 0.3

        # === 目标图例 ===
        ax.text(0.5, y_pos, '【搜救目标】', fontsize=10, fontweight='bold', color='#2C3E50')
        y_pos -= 0.8

        # 普通目标
        ax.plot(1, y_pos, 'X', markersize=8, color=self.colors['target_unfound'],
               markeredgecolor='#2C3E50', markeredgewidth=1.0)
        ax.text(1.8, y_pos, '普通目标', fontsize=7.5, va='center')
        y_pos -= 0.5

        # 遮蔽目标
        ax.plot(1, y_pos, 's', markersize=8, color=self.colors['target_unfound'],
               markeredgecolor='#E67E22', markeredgewidth=1.5)
        ax.text(1.8, y_pos, '遮蔽目标[遮]', fontsize=7.5, va='center', color='#E67E22')
        y_pos -= 0.5

        # 被困建筑目标
        ax.plot(1, y_pos, 'D', markersize=8, color=self.colors['target_unfound'],
               markeredgecolor='#C0392B', markeredgewidth=1.5)
        ax.text(1.8, y_pos, '被困建筑[困]', fontsize=7.5, va='center', color='#C0392B')
        y_pos -= 0.5
        
        # 移动目标
        ax.plot(1, y_pos, 'o', markersize=8, color=self.colors['target_unfound'],
               markeredgecolor='#8E44AD', markeredgewidth=1.5)
        ax.text(1.8, y_pos, '移动目标[移]', fontsize=7.5, va='center', color='#8E44AD')
        y_pos -= 0.5

        # 已发现
        ax.plot(1, y_pos, '*', markersize=10, color=self.colors['target_found'],
               markeredgecolor='#2C3E50', markeredgewidth=1.0)
        ax.text(1.8, y_pos, '已发现', fontsize=7.5, va='center')
        y_pos -= 0.6

        # === 区域图例 ===
        ax.text(0.5, y_pos, '【搜索区域】', fontsize=10, fontweight='bold', color='#2C3E50')
        y_pos -= 0.8

        zone_legend = [
            (self.colors['zone_unassigned'], '待分配'),
            (self.colors['zone_assigned'], '已分配'),
            (self.colors['zone_completed'], '已完成'),
        ]

        for i, (color, label) in enumerate(zone_legend):
            rect = patches.Rectangle((1 + i*3, y_pos - 0.15), 0.4, 0.3,
                                     facecolor=color, edgecolor='#333', linewidth=1)
            ax.add_patch(rect)
            ax.text(1.8 + i*3, y_pos, label, fontsize=7.5, va='center')
        
        y_pos -= 0.6

        # === 其他元素 ===
        ax.text(0.5, y_pos, '【其他元素】', fontsize=10, fontweight='bold', color='#2C3E50')
        y_pos -= 0.8

        # 基地 & 障碍物 (并排)
        base = patches.RegularPolygon((1, y_pos), numVertices=6, radius=0.2,
                                      facecolor=self.colors['base'], edgecolor='#D68910')
        ax.add_patch(base)
        ax.text(1.8, y_pos, '基地', fontsize=7.5, va='center')

        obs = patches.Circle((5.5, y_pos), 0.2, facecolor=self.colors['obstacle'],
                            edgecolor='#2C3E50', hatch='///')
        ax.add_patch(obs)
        ax.text(6.3, y_pos, '障碍物', fontsize=7.5, va='center')
        
        y_pos -= 0.5

        # 感知范围 & 轨迹 (并排)
        ax.plot(1, y_pos, 'o', markersize=12, markerfacecolor='none',
               markeredgecolor='#3498DB', linestyle=':', markeredgewidth=1.0)
        ax.text(1.8, y_pos, '感知范围', fontsize=7.5, va='center')

        ax.plot([5.2, 5.8], [y_pos, y_pos], '-', color='#E74C3C', linewidth=2, alpha=0.5)
        ax.text(6.3, y_pos, '飞行轨迹', fontsize=7.5, va='center')

        y_pos -= 0.5

        # 统计信息框 - 调整字体和位置
        self.stats_box = ax.text(
            5, 0.5, '', fontsize=9, # 字体调小
            ha='center', va='bottom', # 底部对齐
            bbox=dict(boxstyle='round,pad=0.4', facecolor='#EBF5FB',
                     edgecolor='#3498DB', linewidth=1.5)
        )

    def update(self, drones: List[DroneAgent], coordinator: CoordinatorAgent):
        """更新可视化"""# 更新可视化
        
        # 更新覆盖热力图
        if self.coverage_layer:
            grid_data = np.array(self.env.coverage_map.grid)
            self.coverage_layer.set_data(grid_data)

        # 更新无人机
        for drone in drones:
            if drone.agent_id in self.drone_plots:
                plot, idx = self.drone_plots[drone.agent_id]
                color = self.drone_colors.get(idx, '#333333')

                # 更新位置
                plot.set_data([drone.position.x], [drone.position.y])

                # 更新标记形状
                marker = self.state_markers.get(drone.drone_state, '^')
                plot.set_marker(marker)

                # 更新感知范围
                if drone.agent_id in self.drone_ranges:
                    self.drone_ranges[drone.agent_id].center = (drone.position.x, drone.position.y)

                # 更新标签 - 紧贴无人机
                if drone.agent_id in self.drone_labels:
                    label = self.drone_labels[drone.agent_id]
                    label.xy = (drone.position.x, drone.position.y)
                    role = getattr(drone, 'current_role', None)
                    role_str = f"[{role.value}]" if role else ""
                    label.set_text(f'{drone.name}{role_str}\n{drone.drone_state.value}\n电量:{drone.battery:.0f}%')

                # 更新轨迹 - 如果是生成报告模式，保留所有点；否则只保留最近30个点
                trail, positions = self.trail_plots[drone.agent_id]
                positions.append((drone.position.x, drone.position.y))
                
                # 可通过 max_trail_length 控制，设置为 None 或 <=0 表示无限长
                max_trail_length = getattr(self, 'max_trail_length', 30)
                if max_trail_length > 0 and len(positions) > max_trail_length:  
                    positions.pop(0)
                    
                if positions:
                    xs, ys = zip(*positions)
                    trail.set_data(xs, ys)

        # 更新目标状态
        for target in self.env.targets:
            if target.target_id in self.target_plots:
                plot = self.target_plots[target.target_id]
                
                # 更新位置（针对移动目标）
                plot.set_data([target.position.x], [target.position.y])
                
                if target.found:
                    plot.set_color(self.colors['target_found'])
                    plot.set_marker('*')
                    plot.set_markersize(20)

        # 更新区域状态
        for zone in self.env.search_zones:
            if zone.zone_id in self.zone_patches:
                rect = self.zone_patches[zone.zone_id]
                if zone.searched:
                    rect.set_facecolor(self.colors['zone_completed'])
                elif zone.assigned_to:
                    rect.set_facecolor(self.colors['zone_assigned'])

        # 更新统计信息
        stats = self.env.get_statistics()
        status = coordinator.get_global_status()

        # 计算角色分布
        role_counts = {}
        for drone in drones:
            role = getattr(drone, 'current_role', None)
            if role:
                role_counts[role.value] = role_counts.get(role.value, 0) + 1
        role_str = ", ".join([f"{k}:{v}" for k, v in role_counts.items()]) if role_counts else "N/A"

        # 获取算法统计
        algo_stats = getattr(coordinator, 'algorithm_stats', {})
        entropy_str = ""
        if coordinator.swarm_coordinator:
            try:
                entropy = coordinator.swarm_coordinator.info_search.compute_entropy()
                entropy_str = f"Entropy: {entropy:.1f}\n"
            except:
                pass

        stats_text = (
            f"=== Mission Stats ===\n\n"
            f"Time: {stats['time_step']} steps\n\n"
            f"Targets: {stats['found_targets']}/{stats['total_targets']}\n"
            f"Progress: {stats['search_progress']*100:.1f}%\n\n"
            f"Coverage: {stats.get('coverage_ratio', 0)*100:.1f}%\n"
            f"Hotspots: {stats.get('hotspots', 0)}\n\n"
            f"Zones:\n"
            f"  Pending: {status['pending_zones']}\n"
            f"  Done: {status['completed_zones']}\n\n"
            f"Drones: {status['total_drones']}\n\n"
            f"=== Advanced Algo ===\n"
            f"{entropy_str}"
            f"Roles: {role_str}\n"
            f"Role Chg: {algo_stats.get('role_changes', 0)}\n"
            f"Predictions: {algo_stats.get('predictions_made', 0)}"
        )
        self.stats_box.set_text(stats_text)

        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def show(self):
        """显示图形"""
        plt.show()

    def save_frame(self, filename: str):
        """保存当前帧"""
        self.fig.savefig(filename, dpi=150, bbox_inches='tight')


def run_visualization(
        environment: Environment,
        coordinator: CoordinatorAgent,
        drones: List[DroneAgent],
        max_steps: int = 200,
        delay: float = 0.1
):
    """
    运行可视化模拟
    """
    plt.ion()

    viz = Visualizer(environment)
    viz.setup(drones)

    # 启动任务
    coordinator.start_mission(environment)

    step = 0
    mission_complete_step = None  # 记录任务完成的步骤
    return_steps = 50  # 任务完成后给无人机返航的步数

    while step < max_steps:
        # 检查是否应该结束模拟
        if mission_complete_step is not None:
            # 任务已完成，检查是否所有无人机都到达基地或超时
            all_at_base = all(
                drone.position.distance_to(environment.base_position) < 5
                for drone in drones
            )
            if all_at_base:
                print(f"\n所有无人机已返回基地!")
                break
            if (step - mission_complete_step) > return_steps:
                print(f"\n返航超时，部分无人机仍在途中")
                break

        # 更新所有智能体
        for drone in drones:
            drone.update(environment)

            # 检查是否发现目标
            perception = drone.perceive(environment)
            for target in perception.get('visible_targets', []):
                environment.mark_target_found(target.target_id, drone.agent_id)

        # 更新指挥中心
        coordinator_perception = coordinator.perceive(environment)
        action = coordinator.decide(coordinator_perception)
        if action:
            coordinator.act(action)

        # 环境步进
        environment.step()

        # 更新可视化
        viz.update(drones, coordinator)
        plt.pause(delay)

        step += 1

        # 检查是否所有目标都找到
        stats = environment.get_statistics()
        if stats['found_targets'] == stats['total_targets'] and mission_complete_step is None:
            print(f"\n[OK] 所有目标已找到! 用时 {step} 步")
            print("[UAV] 无人机正在返回基地...")
            mission_complete_step = step

    plt.ioff()

    # 保存最终结果
    os.makedirs(FIGURES_DIR, exist_ok=True)
    output_file = os.path.join(FIGURES_DIR, "simulation_result.png")
    viz.save_frame(output_file)
    print(f"[SAVE] 结果已保存至 {output_file}")

    viz.show()