import sys
import numpy as np
import rclpy
from rclpy.executors import SingleThreadedExecutor
from PyQt5 import QtWidgets
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtWidgets import QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QCheckBox, QWidget, QMainWindow, QComboBox
from std_msgs.msg import Float64
from mavros_msgs.srv import SetMode, CommandBool, CommandTOL
from mavros_msgs.msg import State, GPSRAW, VfrHud
from sensor_msgs.msg import NavSatFix, BatteryState, Imu
from geometry_msgs.msg import PoseStamped
from mavros.base import SENSOR_QOS

class ROS2Worker(QThread):
    state_signal = pyqtSignal(str, str)  # 新增 namespace 參數
    battery_signal = pyqtSignal(str, float)  # 新增 namespace 參數
    gps_signal = pyqtSignal(str, float, float)  # 新增 namespace 參數
    gps_status_signal = pyqtSignal(str, int, int)
    local_position_signal = pyqtSignal(str, float, float, float)  # 新增 namespace 參數
    imu_signal = pyqtSignal(str, float, float, float, float)
    hdg_signal = pyqtSignal(str, float)
    groundspeed_signal = pyqtSignal(str, float)  # 新增 namespace 與 groundspeed 值

    def __init__(self, namespaces):
        super().__init__()
        namespaces = ['uav1', 'uav2', 'uav3']
        self.node = rclpy.create_node('multi_uav_control_app')
        self.executor = SingleThreadedExecutor()  # 創建單執行器
        self.executor.add_node(self.node)  # 添加節點到執行器
        self.running = True

        for namespace in namespaces:
            self.node.create_subscription(State, f'/{namespace}/state',
                                          lambda msg, ns=namespace: self.state_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(BatteryState, f'/{namespace}/battery',
                                          lambda msg, ns=namespace: self.battery_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(NavSatFix, f'/{namespace}/global_position/global',
                                          lambda msg, ns=namespace: self.gps_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(GPSRAW, f'/{namespace}/gpsstatus/gps1/raw',
                                          lambda msg, ns=namespace: self.gps_status_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(PoseStamped, f'/{namespace}/local_position/pose',
                                          lambda msg, ns=namespace: self.local_position_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(Imu, f'/{namespace}/imu/data', 
                                          lambda msg, ns=namespace: self.imu_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(Float64, f'/{namespace}/global_position/compass_hdg', 
                                          lambda msg, ns=namespace: self.heading_callback(msg, ns), qos_profile=SENSOR_QOS)
            self.node.create_subscription(VfrHud, f'/{namespace}/vfr_hud',
                                        lambda msg, ns=namespace: self.vfr_hud_callback(msg, ns), qos_profile=SENSOR_QOS)

        self.clients = {namespace: {
            "set_mode": self.node.create_client(SetMode, f'/{namespace}/set_mode'),
            "arming": self.node.create_client(CommandBool, f'/{namespace}/cmd/arming'),
            "takeoff": self.node.create_client(CommandTOL, f'/{namespace}/cmd/takeoff'),
            "land": self.node.create_client(CommandTOL, f'/{namespace}/cmd/land'),
            "local_position": self.node.create_publisher(PoseStamped, f'/{namespace}/setpoint_position/local', 10)
        } for namespace in namespaces}

    def run(self):
        while self.running and rclpy.ok():
            self.executor.spin_once(timeout_sec=0.1)

    def stop(self):
        self.running = False
        self.executor.shutdown()
        self.node.destroy_node()

    def state_callback(self, msg, namespace):
        self.state_signal.emit(namespace, msg.mode)

    def battery_callback(self, msg, namespace):
        self.battery_signal.emit(namespace, msg.percentage)

    def gps_callback(self, msg, namespace):
        self.gps_signal.emit(namespace, msg.latitude, msg.longitude)

    def gps_status_callback(self, msg, namespace):
        self.gps_status_signal.emit(namespace, msg.fix_type, msg.satellites_visible)

    def local_position_callback(self, msg, namespace):
        self.local_position_signal.emit(namespace, msg.pose.position.x, msg.pose.position.y, msg.pose.position.z)

    def imu_callback(self, msg, namespace):
        self.imu_signal.emit(namespace, msg.orientation.x, msg.orientation.y, msg.orientation.z, msg.orientation.w)

    def heading_callback(self, msg, namespace):
        self.hdg_signal.emit(namespace, msg.data)

    def vfr_hud_callback(self, msg, namespace):
        self.groundspeed_signal.emit(namespace, msg.groundspeed)
    
    def set_mode(self, namespace, mode):
        req = SetMode.Request()
        req.custom_mode = mode
        self.clients[namespace]["set_mode"].call_async(req)
        self.node.get_logger().info(f'{namespace}: 模式切換: {mode}')

    def arm(self, namespace, arm: bool):
        req = CommandBool.Request()
        req.value = arm
        self.clients[namespace]["arming"].call_async(req)
        self.node.get_logger().info(f'{namespace}: {"解鎖" if arm else "上鎖"}')

    def takeoff(self, namespace, altitude):
        req = CommandTOL.Request()
        req.altitude = altitude
        req.latitude = 0.0
        req.longitude = 0.0
        req.yaw = 0.0
        req.min_pitch = 0.0
        self.clients[namespace]["takeoff"].call_async(req)
        self.node.get_logger().info(f'{namespace}: 起飛到高度 {altitude}')

    def land(self, namespace):
        req = CommandTOL.Request()
        req.altitude = 0.0
        req.latitude = 0.0
        req.longitude = 0.0
        req.yaw = 0.0
        req.min_pitch = 0.0
        self.clients[namespace]["land"].call_async(req)
        self.node.get_logger().info(f'{namespace}: 降落')

    def set_local_position(self, namespace, x, y, z):
        pose_msg = PoseStamped()
        pose_msg.pose.position.x = x
        pose_msg.pose.position.y = y
        pose_msg.pose.position.z = z
        self.clients[namespace]["local_position"].publish(pose_msg)
        self.node.get_logger().info(f'{namespace}: Published local position: X={x}, Y={y}, Z={z}')

class DroneControlApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.namespaces = ['uav1', 'uav2', 'uav3']
        self.workers = ROS2Worker(self.namespaces)
        self.initUI()
        # 訂閱所有數據更新
        self.workers.state_signal.connect(self.update_state)
        self.workers.battery_signal.connect(self.update_battery)
        self.workers.gps_signal.connect(self.update_gps)
        self.workers.gps_status_signal.connect(self.update_gps_status)
        self.workers.local_position_signal.connect(self.update_local_position)
        self.workers.imu_signal.connect(self.update_imu)
        self.workers.hdg_signal.connect(self.update_hdg)
        self.workers.groundspeed_signal.connect(self.update_groundspeed)
        self.workers.start()
            
    def initUI(self):
        self.setWindowTitle("多無人機控制介面")
        self.setGeometry(100, 100, 800, 600)

        # 主佈局
        main_layout = QVBoxLayout()

        # 無人機選擇區域
        uav_layout = QHBoxLayout()
        self.uav_checkboxes = {}
        for namespace in self.namespaces:
            checkbox = QCheckBox(namespace)
            checkbox.setChecked(True)  # 預設勾選
            self.uav_checkboxes[namespace] = checkbox
            uav_layout.addWidget(checkbox)
        main_layout.addLayout(uav_layout)

        # 顯示所有無人機資訊，從左到右顯示
        uav_layout = QHBoxLayout()
        
        # 逐個顯示 UAV 的資訊
        self.uav_labels = {}
        for namespace in self.namespaces:
            uav_info_layout = QVBoxLayout()  # 每台 UAV 的資訊垂直排列

            # 狀態顯示
            self.uav_labels[namespace] = {
                "status": QLabel(f"{namespace} 狀態：等待數據..."),
                "battery": QLabel(f"{namespace} 電量：等待數據..."),
                "gps": QLabel(f"{namespace} GPS位置：等待數據...\n\n"),
                "fix": QLabel(f"{namespace} Fix Type：等待數據..."),
                "satellites": QLabel(f"{namespace} 衛星數量：等待數據..."),
                "local_position": QLabel(f"{namespace} Local Position：等待數據..."),
                "groundspeed": QLabel(f"{namespace} 地面速度：等待數據..."),
                "pitch": QLabel(f"{namespace} Pitch：等待數據..."),
                "heading": QLabel(f"{namespace} Heading：等待數據..."),
            }

            # 把每個資訊添加到該 UAV 的垂直佈局中
            for label in self.uav_labels[namespace].values():
                uav_info_layout.addWidget(label)

            # 將該 UAV 的資訊佈局加到主佈局中（從左到右排列）
            uav_layout.addLayout(uav_info_layout)

        main_layout.addLayout(uav_layout)

        # 模式切換區域
        mode_layout = QHBoxLayout()
        self.modeComboBox = QComboBox()
        self.modeComboBox.addItems(["GUIDED", "STABILIZE", "AUTO", "LOITER"])
        self.modeButton = QPushButton("切換模式")
        self.modeButton.clicked.connect(self.change_mode)
        mode_layout.addWidget(QLabel("選擇模式："))
        mode_layout.addWidget(self.modeComboBox)
        mode_layout.addWidget(self.modeButton)
        main_layout.addLayout(mode_layout)

        # XYZ 設置欄位
        xyz_layout = QHBoxLayout()
        self.positionX = QLineEdit()
        self.positionX.setPlaceholderText("X")
        self.positionY = QLineEdit()
        self.positionY.setPlaceholderText("Y")
        self.positionZ = QLineEdit()
        self.positionZ.setPlaceholderText("Z")
        xyz_layout.addWidget(QLabel("輸入位置："))
        xyz_layout.addWidget(self.positionX)
        xyz_layout.addWidget(self.positionY)
        xyz_layout.addWidget(self.positionZ)
        main_layout.addLayout(xyz_layout)

        # 設置 XYZ 的按鈕
        self.setPositionButton = QPushButton("設置位置")
        self.setPositionButton.clicked.connect(self.set_local_position)
        main_layout.addWidget(self.setPositionButton)

        # 解鎖按鈕
        self.armButton = QPushButton("解鎖")
        self.armButton.clicked.connect(self.arm_drone)
        main_layout.addWidget(self.armButton)

        # 起飛按鈕
        self.takeoffButton = QPushButton("起飛")
        self.takeoffButton.clicked.connect(self.takeoff_drone)
        main_layout.addWidget(self.takeoffButton)

        # 降落按鈕
        self.landButton = QPushButton("降落")
        self.landButton.clicked.connect(self.land_drone)
        main_layout.addWidget(self.landButton)

        # 設置主佈局
        central_widget = QWidget(self)
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)
        self.show()

    def closeEvent(self, event):
        try:
            self.workers.stop()
        finally:
            event.accept()

    def update_state(self, namespace, mode):
        self.uav_labels[namespace]["status"].setText(f"{namespace} 狀態：{mode}")

    def update_battery(self, namespace, percentage):
        self.uav_labels[namespace]["battery"].setText(f"{namespace} 電量：{percentage * 100:.2f}%")

    def update_gps(self, namespace, latitude, longitude):
        self.uav_labels[namespace]["gps"].setText(f"{namespace} GPS位置： \nLat:{latitude:.6f}° \nLon:{longitude:.6f}°")

    def update_gps_status(self, namespace, fix_type, satellites_visible):
        fix_type_str = {
            0: "No Fix",
            1: "No Fix",
            2: "2D Fix",
            3: "3D Fix",
            4: "RTK Float",
            5: "RTK Fix"
        }.get(fix_type, "Unknown")
        self.uav_labels[namespace]["fix"].setText(f"{namespace} Fix Type：{fix_type_str}")
        self.uav_labels[namespace]["satellites"].setText(f"{namespace} 衛星數量：{satellites_visible}")

    def update_local_position(self, namespace, x, y, z):
        self.uav_labels[namespace]["local_position"].setText(f"{namespace} Local：({x:.2f}, {y:.2f}, {z:.2f})")

    def update_groundspeed(self, namespace, groundspeed):
        self.uav_labels[namespace]["groundspeed"].setText(f"{namespace} 地面速度：{groundspeed:.2f} m/s")

    def update_imu(self, namespace, x, y, z, w):
        q = np.array([x, y, z, w])
        pitch = self.quaternion_to_euler(q)
        self.uav_labels[namespace]["pitch"].setText(f"{namespace} Pitch：{pitch:.2f}°")

    def update_hdg(self, namespace, data):
        self.uav_labels[namespace]["heading"].setText(f"{namespace} Heading：{data:.2f}°")

    def change_mode(self):
        try:
            selected_mode = self.modeComboBox.currentText()
            for namespace in self.namespaces:
                if self.uav_checkboxes[namespace].isChecked():
                    self.workers.set_mode(namespace, selected_mode)
        except Exception as e:
            QtWidgets.QMessageBox.warning(self, "錯誤", f"模式切換失敗：{str(e)}")

    def set_local_position(self):
        try:
            x = float(self.positionX.text().strip())
            y = float(self.positionY.text().strip())
            z = float(self.positionZ.text().strip())
            for namespace in self.namespaces:
                if self.uav_checkboxes[namespace].isChecked():
                    self.workers.set_local_position(namespace, x, y, z)
        except ValueError:
            QtWidgets.QMessageBox.warning(self, "錯誤", "請輸入有效的數值！")

    def arm_drone(self):
        try:
            for namespace in self.namespaces:
                if self.uav_checkboxes[namespace].isChecked():
                    self.workers.arm(namespace, True)
        except Exception as e:
            QtWidgets.QMessageBox.warning(self, "錯誤", f"解鎖失敗：{e}")

    def takeoff_drone(self):
        try:
            z_text = self.positionZ.text().strip()
            altitude = float(z_text) if z_text else 5.0
            for namespace in self.namespaces:
                if self.uav_checkboxes[namespace].isChecked():
                    self.workers.takeoff(namespace, altitude)
        except ValueError:
            QtWidgets.QMessageBox.warning(self, "錯誤", "請輸入有效的數值！")

    def land_drone(self):
        for namespace in self.namespaces:
            if self.uav_checkboxes[namespace].isChecked():
                self.workers.land(namespace)

    def quaternion_to_euler(self, q):
        # Convert quaternion to Euler angles (roll, pitch, yaw)
        x, y, z, w = q

        # Pitch (y-axis rotation)
        pitch_y = np.arcsin(2.0 * (w * y - z * x))
        
        return pitch_y
        
def main(args=None):
    rclpy.init(args=args)
    app = QtWidgets.QApplication(sys.argv)
    window = DroneControlApp()
    window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()