#!/usr/bin/env python3
"""
导航模拟节点 (Nav Simulator Node)
用于模拟 TF 变换和 Nav2 Action Server，以便在没有真实机器人/定位系统时测试导航功能。

功能:
1. 发布 TF 变换 (map → base_footprint)
2. 提供 navigate_to_pose Action Server
3. 模拟机器人移动过程
4. 加载 map_mapping.yaml 中的导航点配置
"""

import rclpy
from rclpy.node import Node
from rclpy.action import ActionServer
from rclpy.action.server import ServerGoalHandle
from rclpy.callback_groups import ReentrantCallbackGroup

import yaml
import math
import time
import os

from geometry_msgs.msg import PoseStamped, TransformStamped, Quaternion
from nav2_msgs.action import NavigateToPose
from tf2_ros import TransformBroadcaster


def quaternion_from_euler(roll, pitch, yaw):
    """从欧拉角 (roll, pitch, yaw) 计算四元数"""
    cy = math.cos(yaw * 0.5)
    sy = math.sin(yaw * 0.5)
    cp = math.cos(pitch * 0.5)
    sp = math.sin(pitch * 0.5)
    cr = math.cos(roll * 0.5)
    sr = math.sin(roll * 0.5)

    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy

    return (x, y, z, w)


class NavSimulatorNode(Node):
    """
    导航模拟节点

    模拟:
    - TF 变换发布 (map → base_footprint)
    - Nav2 NavigateToPose Action Server
    - 机器人移动过程
    """

    def __init__(self):
        super().__init__('nav_simulator_node')

        # ========== 参数声明 ==========
        self.declare_parameter('map_mapping_file', '/home/sunrise/opt/dev/project/aiagent/brain/PlannerNode2/largemodel/config/map_mapping.yaml')
        self.declare_parameter('initial_x', 0.0)
        self.declare_parameter('initial_y', 0.0)
        self.declare_parameter('initial_z', 0.0)
        self.declare_parameter('initial_yaw', 0.0)
        self.declare_parameter('move_speed', 0.5)  # m/s
        self.declare_parameter('tf_publish_rate', 10.0)  # Hz
        self.declare_parameter('simulate_move', True)  # 是否模拟移动过程

        self.map_mapping_file = self.get_parameter('map_mapping_file').value
        self.move_speed = self.get_parameter('move_speed').value
        self.tf_publish_rate = self.get_parameter('tf_publish_rate').value
        self.simulate_move = self.get_parameter('simulate_move').value

        # ========== 状态变量 ==========
        self.navpose_dict = {}  # 导航点字典
        self.current_pose = {
            'x': self.get_parameter('initial_x').value,
            'y': self.get_parameter('initial_y').value,
            'z': self.get_parameter('initial_z').value,
            'yaw': self.get_parameter('initial_yaw').value,
        }
        self.navigation_in_progress = False
        self.cancel_requested = False

        # ========== 加载地图配置 ==========
        self.load_map_mapping()

        # ========== TF 发布者 ==========
        self.tf_broadcaster = TransformBroadcaster(self)
        self.tf_timer = self.create_timer(
            1.0 / self.tf_publish_rate,
            self.publish_tf
        )

        # ========== Nav2 Action Server ==========
        self.action_server = ActionServer(
            self,
            NavigateToPose,
            'navigate_to_pose',
            self.execute_callback,
            callback_group=ReentrantCallbackGroup()
        )

        self.get_logger().info('导航模拟节点已启动')
        self.get_logger().info(f'初始位置: x={self.current_pose["x"]:.3f}, y={self.current_pose["y"]:.3f}, yaw={self.current_pose["yaw"]:.3f}')
        self.get_logger().info(f'移动速度: {self.move_speed} m/s')
        self.get_logger().info(f'已加载 {len(self.navpose_dict)} 个导航点')

    def load_map_mapping(self):
        """从 YAML 文件加载地图导航点配置"""
        if not os.path.exists(self.map_mapping_file):
            self.get_logger().warn(f'地图配置文件不存在: {self.map_mapping_file}')
            return

        try:
            with open(self.map_mapping_file, 'r') as f:
                target_points = yaml.safe_load(f)

            for name, data in target_points.items():
                pose = PoseStamped()
                pose.header.frame_id = 'map'
                pose.pose.position.x = data['position']['x']
                pose.pose.position.y = data['position']['y']
                pose.pose.position.z = data['position'].get('z', 0.0)

                q = quaternion_from_euler(
                    0.0,  # roll
                    0.0,  # pitch
                    data['orientation'].get('yaw', 0.0)
                )
                pose.pose.orientation.x = q[0]
                pose.pose.orientation.y = q[1]
                pose.pose.orientation.z = q[2]
                pose.pose.orientation.w = q[3]

                self.navpose_dict[name] = pose

            self.get_logger().info(f'成功加载地图配置: {self.map_mapping_file}')

        except Exception as e:
            self.get_logger().error(f'加载地图配置失败: {e}')

    def quaternion_to_yaw(self, q):
        """从四元数提取 yaw 角"""
        return math.atan2(
            2.0 * (q.w * q.z + q.x * q.y),
            1.0 - 2.0 * (q.y * q.y + q.z * q.z)
        )

    def yaw_to_quaternion(self, yaw):
        """从 yaw 角生成四元数 (绕 Z 轴旋转)"""
        q = quaternion_from_euler(0.0, 0.0, yaw)
        q_msg = Quaternion()
        q_msg.x = q[0]
        q_msg.y = q[1]
        q_msg.z = q[2]
        q_msg.w = q[3]
        return q_msg

    def publish_tf(self):
        """发布 TF 变换 (map → base_footprint)"""
        t = TransformStamped()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = 'map'
        t.child_frame_id = 'base_footprint'

        t.transform.translation.x = self.current_pose['x']
        t.transform.translation.y = self.current_pose['y']
        t.transform.translation.z = self.current_pose['z']

        q = self.yaw_to_quaternion(self.current_pose['yaw'])
        t.transform.rotation.x = q.x
        t.transform.rotation.y = q.y
        t.transform.rotation.z = q.z
        t.transform.rotation.w = q.w

        self.tf_broadcaster.sendTransform(t)

    def calculate_distance(self, x1, y1, x2, y2):
        """计算两点之间的距离"""
        return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)

    def simulate_movement(self, start_x, start_y, start_yaw, target_x, target_y, target_yaw):
        """
        模拟机器人移动到目标位置

        使用线性插值模拟移动过程:
        1. 计算距离和角度
        2. 按固定速度逐步更新位置
        3. 支持取消导航
        """
        distance = self.calculate_distance(start_x, start_y, target_x, target_y)
        move_time = distance / self.move_speed if self.move_speed > 0 else 0.0

        self.get_logger().info(f'开始模拟移动: 距离={distance:.3f}m, 预计时间={move_time:.2f}s')

        if not self.simulate_move or move_time == 0.0:
            # 立即到达目标
            self.current_pose['x'] = target_x
            self.current_pose['y'] = target_y
            self.current_pose['yaw'] = target_yaw
            return True

        # 分步移动，频率为 tf_publish_rate
        steps = max(10, int(move_time * self.tf_publish_rate))
        dt = move_time / steps if steps > 0 else 0.0

        for i in range(steps):
            if self.cancel_requested:
                self.get_logger().info('导航被取消')
                self.cancel_requested = False
                return False

            alpha = (i + 1) / steps
            self.current_pose['x'] = start_x + (target_x - start_x) * alpha
            self.current_pose['y'] = start_y + (target_y - start_y) * alpha

            # 角度插值（考虑角度跨越问题）
            yaw_diff = target_yaw - start_yaw
            while yaw_diff > math.pi:
                yaw_diff -= 2 * math.pi
            while yaw_diff < -math.pi:
                yaw_diff += 2 * math.pi
            self.current_pose['yaw'] = start_yaw + yaw_diff * alpha

            time.sleep(dt)

        # 确保到达精确目标位置
        self.current_pose['x'] = target_x
        self.current_pose['y'] = target_y
        self.current_pose['yaw'] = target_yaw

        return True

    async def execute_callback(self, goal_handle: ServerGoalHandle):
        """
        NavigateToPose Action Server 的执行回调

        处理导航请求:
        1. 接收目标点
        2. 模拟移动过程
        3. 返回结果
        """
        self.navigation_in_progress = True
        self.cancel_requested = False

        goal_pose = goal_handle.request.pose
        target_x = goal_pose.pose.position.x
        target_y = goal_pose.pose.position.y
        target_z = goal_pose.pose.position.z
        target_yaw = self.quaternion_to_yaw(goal_pose.pose.orientation)

        start_x = self.current_pose['x']
        start_y = self.current_pose['y']
        start_yaw = self.current_pose['yaw']

        self.get_logger().info(f'收到导航目标: x={target_x:.3f}, y={target_y:.3f}, yaw={target_yaw:.3f}')
        self.get_logger().info(f'当前位置: x={start_x:.3f}, y={start_y:.3f}, yaw={start_yaw:.3f}')

        # 发布反馈
        feedback_msg = NavigateToPose.Feedback()
        feedback_msg.current_pose = goal_pose
        goal_handle.publish_feedback(feedback_msg)

        # 模拟移动
        success = self.simulate_movement(start_x, start_y, start_yaw, target_x, target_y, target_yaw)

        self.navigation_in_progress = False

        if success:
            self.get_logger().info('导航成功完成')
            goal_handle.succeed()
            return NavigateToPose.Result()
        else:
            self.get_logger().info('导航被取消')
            goal_handle.aborted()
            return NavigateToPose.Result()

    def set_cancel(self):
        """请求取消当前导航"""
        if self.navigation_in_progress:
            self.get_logger().info('收到取消请求')
            self.cancel_requested = True

    def set_position(self, x, y, yaw):
        """
        手动设置机器人位置（用于测试）

        Args:
            x, y: 位置坐标
            yaw: 朝向角 (弧度)
        """
        self.current_pose['x'] = x
        self.current_pose['y'] = y
        self.current_pose['yaw'] = yaw
        self.get_logger().info(f'位置已更新: x={x:.3f}, y={y:.3f}, yaw={yaw:.3f}')


def main(args=None):
    rclpy.init(args=args)

    node = NavSimulatorNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
