import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import random
from collections import deque, namedtuple
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from tqdm import tqdm
import os
import matplotlib.font_manager as fm

# 设置随机种子以确保结果可复现
np.random.seed(42)
torch.manual_seed(42)
random.seed(42)

# 优化图形布局参数
plt.rcParams["figure.figsize"] = (20, 20)  # 增大图形尺寸
plt.rcParams["font.size"] = 8  # 减小默认字体大小
plt.rcParams["axes.titlesize"] = 10  # 标题字体大小
plt.rcParams["axes.labelsize"] = 9  # 坐标轴标签字体大小
plt.rcParams["xtick.labelsize"] = 8  # x轴刻度字体大小
plt.rcParams["ytick.labelsize"] = 8  # y轴刻度字体大小
plt.rcParams["legend.fontsize"] = 8  # 图例字体大小
plt.rcParams["figure.titlesize"] = 14  # 全局标题字体大小


# 自动检测系统可用的中文字体
def get_available_chinese_fonts():
    """获取系统中可用的中文字体列表"""
    chinese_fonts = []
    for font in fm.findSystemFonts():
        try:
            font_prop = fm.FontProperties(fname=font)
            font_name = font_prop.get_name()
            # 检查字体是否支持中文字符
            if font_prop.get_style() != 'unknown' and ('hei' in font_name.lower() or
                                                       'song' in font_name.lower() or
                                                       'kai' in font_name.lower()):
                chinese_fonts.append(font_name)
        except:
            continue
    return chinese_fonts


# 获取可用中文字体
available_fonts = get_available_chinese_fonts()
print(f"可用的中文字体: {available_fonts}")

# 设置中文字体
if available_fonts:
    plt.rcParams["font.family"] = available_fonts[0]
    print(f"使用字体: {available_fonts[0]}")
else:
    print("未找到可用的中文字体，将使用默认字体")
    plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "sans-serif"]

plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# 定义动作空间
class Action:
    STAY = 0
    FORWARD = 1
    BACKWARD = 2
    UP = 3
    DOWN = 4
    LEFT = 5
    RIGHT = 6

    @staticmethod
    def get_action_name(action):
        names = {
            Action.STAY: "停留",
            Action.FORWARD: "前进",
            Action.BACKWARD: "后退",
            Action.UP: "上浮",
            Action.DOWN: "下潜",
            Action.LEFT: "左转",
            Action.RIGHT: "右转"
        }
        return names.get(action, "未知")


# 定义三维海洋环境
class OceanEnvironment:
    def __init__(self, grid_size=(30, 30, 15), obstacle_density=0.15, current_strength=0.2):
        """初始化海洋环境"""
        self.grid_size = grid_size  # (x, y, z)
        self.obstacle_density = obstacle_density
        self.current_strength = current_strength
        self.obstacles = None
        self.currents = None
        self.reset()

    def reset(self):
        """重置环境，生成新的障碍物和洋流分布"""
        # 生成随机障碍物
        self.obstacles = np.random.choice(
            [0, 1],
            size=self.grid_size,
            p=[1 - self.obstacle_density, self.obstacle_density]
        )

        # 设置起点和终点
        self.start_pos = (0, 0, self.grid_size[2] // 2)
        self.target_pos = (self.grid_size[0] - 1, self.grid_size[1] - 1, self.grid_size[2] // 2)

        # 确保起点和终点没有障碍物
        self.obstacles[self.start_pos] = 0
        self.obstacles[self.target_pos] = 0

        # 生成洋流（水流） - 影响AUV的移动成本
        self.currents = np.random.uniform(-self.current_strength, self.current_strength, size=self.grid_size)

        # 记录AUV当前位置
        self.agent_pos = self.start_pos

        # 记录上一个动作，用于计算转弯成本
        self.last_action = None

        # 记录步数
        self.steps = 0
        self.max_steps = self.grid_size[0] * self.grid_size[1] * 2

        return self.get_observation()

    def get_observation(self):
        """获取当前观察"""
        # 观察范围 - 以AUV为中心的立方体区域
        view_range = 3

        # 创建观察区域
        obs = np.zeros((view_range * 2 + 1, view_range * 2 + 1, view_range * 2 + 1, 3))  # 障碍物、水流、距离

        # 填充观察区域
        for dx in range(-view_range, view_range + 1):
            for dy in range(-view_range, view_range + 1):
                for dz in range(-view_range, view_range + 1):
                    # 计算网格中的实际位置
                    gx = self.agent_pos[0] + dx
                    gy = self.agent_pos[1] + dy
                    gz = self.agent_pos[2] + dz

                    # 检查是否在网格范围内
                    if 0 <= gx < self.grid_size[0] and 0 <= gy < self.grid_size[1] and 0 <= gz < self.grid_size[2]:
                        # 障碍物信息
                        obs[dx + view_range, dy + view_range, dz + view_range, 0] = self.obstacles[gx, gy, gz]

                        # 水流信息
                        obs[dx + view_range, dy + view_range, dz + view_range, 1] = self.currents[gx, gy, gz]

                        # 到目标的相对距离
                        dist_to_target = np.sqrt(
                            (gx - self.target_pos[0]) ** 2 +
                            (gy - self.target_pos[1]) ** 2 +
                            (gz - self.target_pos[2]) ** 2
                        )
                        max_possible_dist = np.sqrt(
                            (self.grid_size[0] - 1) ** 2 +
                            (self.grid_size[1] - 1) ** 2 +
                            (self.grid_size[2] - 1) ** 2
                        )
                        obs[dx + view_range, dy + view_range, dz + view_range, 2] = dist_to_target / max_possible_dist

        # 自身位置相对于目标的归一化距离
        rel_pos = np.array([
            (self.agent_pos[0] - self.target_pos[0]) / self.grid_size[0],
            (self.agent_pos[1] - self.target_pos[1]) / self.grid_size[1],
            (self.agent_pos[2] - self.target_pos[2]) / self.grid_size[2]
        ])

        # 将观察展平
        obs_flat = obs.flatten()
        return np.concatenate([obs_flat, rel_pos])

    def is_valid_position(self, pos):
        """检查位置是否有效（在网格内且无障碍物）"""
        x, y, z = pos
        return (
                0 <= x < self.grid_size[0] and
                0 <= y < self.grid_size[1] and
                0 <= z < self.grid_size[2] and
                self.obstacles[x, y, z] == 0
        )

    def step(self, action):
        """执行动作并返回下一个状态、奖励和是否终止"""
        self.steps += 1

        # 计算新位置
        x, y, z = self.agent_pos
        if action == Action.FORWARD:
            x += 1
        elif action == Action.BACKWARD:
            x -= 1
        elif action == Action.UP:
            z += 1
        elif action == Action.DOWN:
            z -= 1
        elif action == Action.LEFT:
            y -= 1
        elif action == Action.RIGHT:
            y += 1

        new_pos = (x, y, z)

        # 检查是否有效位置
        if self.is_valid_position(new_pos):
            self.agent_pos = new_pos
            collision = False
        else:
            # 如果移动到无效位置，保持原位
            collision = True

        # 计算奖励
        reward = self.calculate_reward(action, collision)

        # 检查是否到达目标
        done = self.agent_pos == self.target_pos

        # 检查是否超时
        if self.steps >= self.max_steps:
            done = True
            reward -= 50  # 超时惩罚

        # 更新上一个动作
        self.last_action = action

        return self.get_observation(), reward, done, {}

    def calculate_reward(self, action, collision):
        """计算奖励函数"""
        reward = 0

        # 基础移动成本
        reward -= 1

        # 碰撞惩罚
        if collision:
            reward -= 100
            return reward

        # 到达目标奖励
        if self.agent_pos == self.target_pos:
            reward += 1000
            return reward

        # 计算到目标的距离
        current_dist = np.sqrt(
            (self.agent_pos[0] - self.target_pos[0]) ** 2 +
            (self.agent_pos[1] - self.target_pos[1]) ** 2 +
            (self.agent_pos[2] - self.target_pos[2]) ** 2
        )

        # 距离减少奖励，增加奖励
        if hasattr(self, 'last_dist'):
            dist_change = self.last_dist - current_dist
            reward += dist_change * 10  # 距离减少奖励

        self.last_dist = current_dist

        # 转弯成本 - 转弯比直行更昂贵
        if self.last_action is not None and action != self.last_action:
            reward -= 2  # 转弯惩罚

        # 水流影响 - 顺流奖励，逆流惩罚
        current_flow = self.currents[self.agent_pos]
        if (action == Action.FORWARD and current_flow > 0) or (action == Action.BACKWARD and current_flow < 0):
            reward += 2  # 顺流奖励
        elif (action == Action.FORWARD and current_flow < 0) or (action == Action.BACKWARD and current_flow > 0):
            reward -= 2  # 逆流惩罚

        # 安全风险 - 接近障碍物惩罚
        safety_distance = 2
        for dx in range(-safety_distance, safety_distance + 1):
            for dy in range(-safety_distance, safety_distance + 1):
                for dz in range(-safety_distance, safety_distance + 1):
                    x, y, z = self.agent_pos[0] + dx, self.agent_pos[1] + dy, self.agent_pos[2] + dz
                    if 0 <= x < self.grid_size[0] and 0 <= y < self.grid_size[1] and 0 <= z < self.grid_size[2]:
                        if self.obstacles[x, y, z] == 1:
                            # 障碍物越近，惩罚越大
                            dist = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
                            if dist > 0:
                                reward -= 5 / dist

        return reward

    def render(self, path=None):
        """可视化环境和路径"""
        fig = plt.figure(figsize=(12, 10))
        ax = fig.add_subplot(111, projection='3d')

        # 绘制障碍物
        obs_x, obs_y, obs_z = np.where(self.obstacles == 1)
        ax.scatter(obs_x, obs_y, obs_z, c='red', marker='s', s=20, label='障碍物')

        # 绘制起点和终点
        ax.scatter(self.start_pos[0], self.start_pos[1], self.start_pos[2], c='green', marker='^', s=100, label='起点')
        ax.scatter(self.target_pos[0], self.target_pos[1], self.target_pos[2], c='blue', marker='*', s=100,
                   label='终点')

        # 绘制路径
        if path is not None:
            path_x = [p[0] for p in path]
            path_y = [p[1] for p in path]
            path_z = [p[2] for p in path]
            ax.plot(path_x, path_y, path_z, c='purple', linewidth=2, label='路径')

        # 设置坐标轴标签
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

        # 设置标题
        ax.set_title('三维海洋环境中的AUV路径规划')

        # 设置图例
        ax.legend()

        # 设置视角
        ax.view_init(elev=30, azim=45)

        plt.tight_layout()
        plt.show()


# 定义DQN网络
class DQN(nn.Module):
    def __init__(self, state_size, action_size):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(state_size, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, action_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        return self.fc4(x)


# 经验回放缓冲区
Experience = namedtuple('Experience', ('state', 'action', 'reward', 'next_state', 'done'))


class ReplayBuffer:
    def __init__(self, capacity):
        self.memory = deque(maxlen=capacity)

    def push(self, *args):
        self.memory.append(Experience(*args))

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


# DQN智能体
class DQNAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size

        # 折扣因子和学习率
        self.gamma = 0.99  # 折扣因子
        self.epsilon = 1.0  # 探索率
        self.epsilon_min = 0.01
        self.epsilon_decay = 0.995
        self.learning_rate = 0.001
        self.batch_size = 64

        # 策略网络和目标网络
        self.policy_net = DQN(state_size, action_size).to(device)
        self.target_net = DQN(state_size, action_size).to(device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        # 优化器和经验回放缓冲区
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=self.learning_rate)
        self.memory = ReplayBuffer(10000)

        # 目标网络更新频率
        self.target_update = 10

    def act(self, state):
        """根据ε-贪心策略选择动作"""
        if np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)

        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            q_values = self.policy_net(state)
        return torch.argmax(q_values, dim=1).item()

    def train(self):
        """训练智能体"""
        if len(self.memory) < self.batch_size:
            return 0

        # 从经验回放缓冲区中采样
        experiences = self.memory.sample(self.batch_size)

        # 优化：预先转换为numpy数组再创建张量
        states = torch.FloatTensor(np.array([e.state for e in experiences])).to(device)
        actions = torch.LongTensor(np.array([e.action for e in experiences])).to(device)
        rewards = torch.FloatTensor(np.array([e.reward for e in experiences])).to(device)
        next_states = torch.FloatTensor(np.array([e.next_state for e in experiences])).to(device)
        dones = torch.FloatTensor(np.array([e.done for e in experiences])).to(device)

        # 计算当前Q值和目标Q值
        current_q = self.policy_net(states).gather(1, actions.unsqueeze(1)).squeeze(1)

        with torch.no_grad():
            next_q = self.target_net(next_states).max(1)[0]
            target_q = rewards + (1 - dones) * self.gamma * next_q

        # 计算损失并优化
        loss = F.mse_loss(current_q, target_q)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)

        return loss.item()

    def update_target_network(self):
        """更新目标网络"""
        self.target_net.load_state_dict(self.policy_net.state_dict())

    def save_model(self, path):
        """保存模型参数"""
        torch.save({
            'policy_net_state_dict': self.policy_net.state_dict(),
            'target_net_state_dict': self.target_net.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'epsilon': self.epsilon
        }, path)

    def load_model(self, path):
        """加载模型参数"""
        checkpoint = torch.load(path)
        self.policy_net.load_state_dict(checkpoint['policy_net_state_dict'])
        self.target_net.load_state_dict(checkpoint['target_net_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.epsilon = checkpoint['epsilon']
        self.policy_net.eval()
        self.target_net.eval()


# 训练函数
def train_agent(env, agent, episodes=1000):
    """训练DQN智能体"""
    scores = []
    avg_scores = []
    losses = []

    for episode in tqdm(range(episodes)):
        state = env.reset()
        score = 0
        episode_loss = []

        done = False
        while not done:
            action = agent.act(state)
            next_state, reward, done, _ = env.step(action)
            agent.memory.push(state, action, reward, next_state, done)
            state = next_state
            score += reward

            loss = agent.train()
            if loss > 0:
                episode_loss.append(loss)

        # 每10个episode更新一次目标网络
        if episode % agent.target_update == 0:
            agent.update_target_network()

        # 记录本轮得分和损失
        scores.append(score)
        avg_score = np.mean(scores[-100:]) if len(scores) >= 100 else np.mean(scores)
        avg_scores.append(avg_score)

        if episode_loss:
            avg_loss = np.mean(episode_loss)
            losses.append(avg_loss)
        else:
            losses.append(0)

        # 打印训练信息
        if (episode + 1) % 10 == 0:
            print(
                f"Episode {episode + 1}/{episodes}, 得分: {score:.2f}, 平均得分: {avg_score:.2f}, 探索率: {agent.epsilon:.4f}")

    # 绘制训练曲线
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(scores, label='得分')
    plt.plot(avg_scores, label='平均得分')
    plt.xlabel('回合数')
    plt.ylabel('得分')
    plt.title('训练得分')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(losses)
    plt.xlabel('回合数')
    plt.ylabel('损失')
    plt.title('训练损失')

    plt.tight_layout()
    plt.show()

    return agent


# 评估函数
def evaluate_agent(env, agent, episodes=100):
    """评估训练好的DQN智能体并收集性能数据"""
    success_count = 0
    step_counts = []
    collision_counts = []
    min_distances = []
    turning_counts = []
    against_current_counts = []
    path_lengths = []
    time_costs = []
    energy_costs = []

    for episode in range(episodes):
        state = env.reset()
        done = False
        path = [env.agent_pos]
        step_count = 0
        collision_count = 0
        turning_count = 0
        against_current_count = 0
        last_action = None
        time_cost = 0
        energy_cost = 0

        while not done:
            # 贪婪策略选择动作
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
            with torch.no_grad():
                q_values = agent.policy_net(state_tensor)
            action = torch.argmax(q_values, dim=1).item()

            # 记录转弯次数
            if last_action is not None and action != last_action:
                turning_count += 1
            last_action = action

            # 记录逆水流移动
            if (action == Action.FORWARD and env.currents[env.agent_pos] < 0) or \
                    (action == Action.BACKWARD and env.currents[env.agent_pos] > 0):
                against_current_count += 1

            # 执行动作
            next_state, reward, done, _ = env.step(action)
            state = next_state
            step_count += 1
            path.append(env.agent_pos)

            # 记录碰撞
            if reward <= -100:  # 碰撞惩罚
                collision_count += 1

            # 计算时间成本 (假设每次动作耗时1单位)
            time_cost += 1

            # 计算能量成本 (基础能耗 + 转弯能耗 + 逆水流能耗)
            base_energy = 1
            turning_energy = 2 if (last_action is not None and action != last_action) else 0
            current_energy = 2 if (action == Action.FORWARD and env.currents[env.agent_pos] < 0) or \
                                  (action == Action.BACKWARD and env.currents[env.agent_pos] > 0) else 0
            energy_cost += base_energy + turning_energy + current_energy

        # 计算与障碍物的最小距离
        min_distance = float('inf')
        for pos in path:
            for dx in range(-2, 3):
                for dy in range(-2, 3):
                    for dz in range(-2, 3):
                        x, y, z = pos[0] + dx, pos[1] + dy, pos[2] + dz
                        if 0 <= x < env.grid_size[0] and 0 <= y < env.grid_size[1] and 0 <= z < env.grid_size[2]:
                            if env.obstacles[x, y, z] == 1:
                                dist = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
                                if dist < min_distance:
                                    min_distance = dist

        # 计算路径长度
        path_length = 0
        for i in range(1, len(path)):
            dx = path[i][0] - path[i - 1][0]
            dy = path[i][1] - path[i - 1][1]
            dz = path[i][2] - path[i - 1][2]
            path_length += np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)

        # 记录结果
        if env.agent_pos == env.target_pos:
            success_count += 1

        step_counts.append(step_count)
        collision_counts.append(collision_count)
        min_distances.append(min_distance)
        turning_counts.append(turning_count)
        against_current_counts.append(against_current_count)
        path_lengths.append(path_length)
        time_costs.append(time_cost)
        energy_costs.append(energy_cost)

    # 计算统计数据
    success_rate = success_count / episodes
    avg_steps = np.mean(step_counts)
    avg_collisions = np.mean(collision_counts)
    avg_min_distance = np.mean(min_distances)
    avg_turning = np.mean(turning_counts)
    avg_against_current = np.mean(against_current_counts)
    avg_path_length = np.mean(path_lengths)
    avg_time_cost = np.mean(time_costs)
    avg_energy_cost = np.mean(energy_costs)

    return {
        'success_rate': success_rate,
        'avg_steps': avg_steps,
        'avg_collisions': avg_collisions,
        'avg_min_distance': avg_min_distance,
        'avg_turning': avg_turning,
        'avg_against_current': avg_against_current,
        'avg_path_length': avg_path_length,
        'avg_time_cost': avg_time_cost,
        'avg_energy_cost': avg_energy_cost,
        'step_counts': step_counts,
        'min_distances': min_distances,
        'turning_counts': turning_counts,
        'against_current_counts': against_current_counts,
        'time_costs': time_costs,
        'energy_costs': energy_costs
    }


# 主函数
def main():
    # 创建环境
    env = OceanEnvironment(grid_size=(30, 30, 15), obstacle_density=0.15, current_strength=0.2)

    # 获取状态和动作空间维度
    state_size = env.get_observation().shape[0]
    action_size = 7

    # 创建智能体
    agent = DQNAgent(state_size, action_size)

    # 模型保存路径
    model_path = 'auv_dqn_model.pth'

    # 检查是否存在预训练模型
    if os.path.exists(model_path):
        user_input = input("发现预训练模型，是否加载模型跳过训练？(y/n): ")
        if user_input.lower() == 'y':
            print("加载预训练模型...")
            agent.load_model(model_path)
        else:
            print("开始训练智能体...")
            trained_agent = train_agent(env, agent, episodes=200)
            print("保存训练好的模型...")
            trained_agent.save_model(model_path)
    else:
        print("未找到预训练模型，开始训练...")
        trained_agent = train_agent(env, agent, episodes=200)
        print("保存训练好的模型...")
        trained_agent.save_model(model_path)

    # 评估智能体
    print("\n开始评估智能体...")
    metrics = evaluate_agent(env, agent, episodes=50)

    # 打印评估结果
    print("\n性能评估结果:")
    print(f"成功率: {metrics['success_rate'] * 100:.2f}%")
    print(f"平均步数: {metrics['avg_steps']:.2f}")
    print(f"平均碰撞次数: {metrics['avg_collisions']:.2f}")
    print(f"与障碍物的平均最小距离: {metrics['avg_min_distance']:.2f}")
    print(f"平均转弯次数: {metrics['avg_turning']:.2f}")
    print(f"平均逆水流移动次数: {metrics['avg_against_current']:.2f}")
    print(f"平均路径长度: {metrics['avg_path_length']:.2f}")
    print(f"平均时间成本: {metrics['avg_time_cost']:.2f}")
    print(f"平均能量成本: {metrics['avg_energy_cost']:.2f}")

    # 可视化评估结果 - 优化标题布局
    fig = plt.figure(figsize=(22, 22))  # 增大图形尺寸

    # 调整子图布局
    plt.subplots_adjust(
        left=0.05,
        right=0.95,
        bottom=0.05,
        top=0.93,  # 减少顶部边距
        wspace=0.3,
        hspace=0.4  # 增加子图之间的垂直间距
    )

    # 使用更紧凑的子图标题
    subplot_titles = [
        '步数分布', '安全距离',
        '转弯次数', '逆水流次数', '能量成本',
        '时间 vs 能量', '安全 vs 效率', '典型路径'
    ]

    # 创建子图
    axes = []
    for i in range(8):
        if i == 7:  # 最后一个子图是3D图
            ax = fig.add_subplot(2, 4, i + 1, projection='3d')
        else:
            ax = fig.add_subplot(2, 4, i + 1)
        axes.append(ax)

    # 1. 步数分布 (时间相关)
    axes[0].hist(metrics['step_counts'], bins=10)
    axes[0].set_title(subplot_titles[0])
    axes[0].set_xlabel('步数')
    axes[0].set_ylabel('频率')

    # 2. 安全性评估
    axes[1].hist(metrics['min_distances'], bins=10)
    axes[1].axvline(x=1.0, color='r', linestyle='--', label='安全阈值')
    axes[1].set_title(subplot_titles[1])
    axes[1].set_xlabel('距离')
    axes[1].set_ylabel('频率')
    axes[1].legend()

    # 3. 能耗评估 - 转弯次数
    axes[2].hist(metrics['turning_counts'], bins=10)
    axes[2].set_title(subplot_titles[2])
    axes[2].set_xlabel('转弯次数')
    axes[2].set_ylabel('频率')

    # 4. 能耗评估 - 逆水流次数
    axes[3].hist(metrics['against_current_counts'], bins=10)
    axes[3].set_title(subplot_titles[3])
    axes[3].set_xlabel('逆水流次数')
    axes[3].set_ylabel('频率')

    # 5. 能量成本分布
    axes[4].hist(metrics['energy_costs'], bins=10)
    axes[4].set_title(subplot_titles[4])
    axes[4].set_xlabel('能量成本')
    axes[4].set_ylabel('频率')

    # 6. 时间与能量成本关系
    axes[5].scatter(metrics['time_costs'], metrics['energy_costs'])
    axes[5].set_title(subplot_titles[5])
    axes[5].set_xlabel('时间成本')
    axes[5].set_ylabel('能量成本')

    # 7. 安全性与效率关系
    axes[6].scatter(metrics['min_distances'], metrics['step_counts'])
    axes[6].set_title(subplot_titles[6])
    axes[6].set_xlabel('最小安全距离')
    axes[6].set_ylabel('步数')

    # 8. 显示一条典型路径
    ax8 = fig.add_subplot(2, 4, 8, projection='3d')

    # 绘制障碍物
    obs_x, obs_y, obs_z = np.where(env.obstacles == 1)
    ax8.scatter(obs_x, obs_y, obs_z, c='red', marker='s', s=20, label='障碍物')

    # 绘制起点和终点
    ax8.scatter(env.start_pos[0], env.start_pos[1], env.start_pos[2], c='green', marker='^', s=100, label='起点')
    ax8.scatter(env.target_pos[0], env.target_pos[1], env.target_pos[2], c='blue', marker='*', s=100, label='终点')

    # 绘制一条典型路径
    state = env.reset()
    done = False
    path = [env.agent_pos]

    while not done:
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            q_values = agent.policy_net(state_tensor)
        action = torch.argmax(q_values, dim=1).item()

        next_state, reward, done, _ = env.step(action)
        state = next_state
        path.append(env.agent_pos)

    path_x = [p[0] for p in path]
    path_y = [p[1] for p in path]
    path_z = [p[2] for p in path]
    ax8.plot(path_x, path_y, path_z, c='purple', linewidth=2, label='路径')

    ax8.set_title('典型AUV路径')
    ax8.set_xlabel('X')
    ax8.set_ylabel('Y')
    ax8.set_zlabel('Z')
    ax8.legend()

    # 调整视角
    ax8.view_init(elev=20, azim=45)

    # 设置总标题，微调位置
    fig.suptitle('AUV路径规划性能评估', fontsize=14, y=0.98)

    plt.show()


if __name__ == "__main__":
    main()