import argparse
import time
from swarmae.SwarmAEClient import SwarmAEClient
import math
from concurrent.futures import ThreadPoolExecutor

class SwarmaeClass:
    def __init__(self):
        self.ae_client = None
        self.nodes = []
        self.node_info = {}  # 存储每个无人机的初始位置和高度
        self.positions = []  # 存储每个无人机的当前位置

    def create_client(self, ue_ip="127.0.0.1", ue_port=2000):
        print("--Connecting to UE4 simulation:", ue_ip, ue_port)
        self.ae_client = SwarmAEClient(ue_ip=ue_ip, ue_port=ue_port)
        return self.ae_client

    def create_node(self, id, max_retries=3):
        for attempt in range(max_retries):
            frame_timestamp = int(round(time.time() * 1000))
            frame_timestamp, sw_node, sw_code = self.ae_client.register_node(
                node_type="四旋翼", node_name="无人机" + str(id), node_no=id, frame_timestamp=frame_timestamp
            )
            
            if sw_code == 200:
                node_name, node_no, team, node_type, _, _ = sw_node.get_node_info()
                initial_location = sw_node.get_location()[:3]
                print(f"--UAV {id} registered successfully: {node_name}, Type: {node_type}")
                print(f"--Initial position: x={initial_location[0]:.2f}, y={initial_location[1]:.2f}, z={initial_location[2]:.2f}")
                self.node_info[sw_node] = {
                    'flight_altitude': initial_location[2] + 1.0,  # 设置飞行高度为起始z位置加1米
                    'ground_height': initial_location[2] + 0.1  # 设置地面高度为初始z位置
                }
                self.positions.append((sw_node, id, initial_location))  # 保存节点及其编号
                self.nodes.append(sw_node)
                return sw_code, sw_node
            else:
                print(f"--UAV {id} registration failed with code {sw_code}, attempt {attempt+1} of {max_retries}")
                time.sleep(2)  # 等待2秒后重试

        print(f"UAV {id} registration failed after {max_retries} attempts.")
        return sw_code, None

    def move_to_position_xy(self, node, target_x, target_y, target_z):
        node_index = self.nodes.index(node)
        current_timestamp = int(round(time.time() * 1000))
        node.control_kinetic_simply_global(x=target_x, y=target_y, z=target_z, v=3.0, frame_timestamp=current_timestamp)
        while True:
            current_location = node.get_location()[:3]
            distance = math.sqrt((current_location[0] - target_x) ** 2 + (current_location[1] - target_y) ** 2)
            self.update_position(node)  # 实时更新位置
            if distance <= 1.5:
                print(f"[UAV {node_index + 1}] reached target position: x={target_x:.2f}, y={target_y:.2f}")
                break
            time.sleep(0.1)

    def move_to_position_z(self, node, target_x, target_y, target_z):
        node_index = self.nodes.index(node)
        current_timestamp = int(round(time.time() * 1000))
        node.control_kinetic_simply_global(x=target_x, y=target_y, z=target_z, v=3.0, frame_timestamp=current_timestamp)
        while True:
            current_location = node.get_location()[:3]
            self.update_position(node)  # 实时更新位置
            if abs(current_location[2] - target_z) <= 0.2:
                print(f"[UAV {node_index + 1}] reached target altitude: z={target_z:.2f}")
                break
            time.sleep(0.1)

    def update_position(self, node):
        """更新并保存无人机的当前位置。"""
        current_location = node.get_location()[:3]
        node_index = self.nodes.index(node)
        self.positions[node_index] = (node, node_index + 1, current_location)

    def takeoff(self, node):
        location = node.get_location()
        x, y, z = location[:3]
        target_z = self.node_info[node]['flight_altitude']  # 获取对应的飞行高度
        self.move_to_position_z(node, x + 0.5, y + 0.5, target_z)  # 垂直起飞需要在x，y上添加一点偏移，否则起飞过慢

    def move_forward(self, node, distance):
        location = node.get_location()
        x, y, z = location[:3]
        target_x = x + distance
        self.move_to_position_xy(node, target_x, y, z)

    def move_backward(self, node, distance):
        location = node.get_location()
        x, y, z = location[:3]
        target_x = x - distance
        self.move_to_position_xy(node, target_x, y, z)

    def move_left(self, node, distance):
        location = node.get_location()
        x, y, z = location[:3]
        target_y = y - distance
        self.move_to_position_xy(node, x, target_y, z)

    def move_right(self, node, distance):
        location = node.get_location()
        x, y, z = location[:3]
        target_y = y + distance
        self.move_to_position_xy(node, x, target_y, z)

    def execute_flight_plan(self, flight_plan_funcs):
        """执行所有无人机的飞行计划。"""
        if len(self.nodes) < len(flight_plan_funcs):
            print("Not all nodes were successfully registered. Adjusting the flight plans.")
            flight_plan_funcs = flight_plan_funcs[:len(self.nodes)]

        with ThreadPoolExecutor() as executor:
            futures = []
            for node, flight_plan_func in zip(self.nodes, flight_plan_funcs):
                futures.append(executor.submit(flight_plan_func, node))
            for future in futures:
                future.result()  # 等待所有无人机完成其飞行计划

if __name__ == '__main__':
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("-i", "--ip", default="127.0.0.1", help="指定服务器IP", type=str)
        parser.add_argument("-p", "--port", default=2000, help="指定服务器Port", type=int)
        parser.add_argument("-n", "--number", default=1, help="指定飞机数量", type=int)
        args = parser.parse_args()

        t1 = SwarmaeClass()
        t1.create_client(args.ip, args.port)
        for i in range(args.number):
            t1.create_node(i+1)
            
        
        #####################
        # 定义飞行计划
        # 前面的部分不需要改，要控制无人机的飞行，只需修改下面的部分即可
        # flight_plan_n即第n个无人机的飞行计划函数
        # t1.takeoff(node)  # 起飞
        # t1.move_forward(node, distance)  # 向前移动distance米
        # t1.move_backward(node, distance)  # 向后移动distance米
        # t1.move_left(node, distance)  # 向左移动distance米
        # t1.move_right(node, distance)  # 向右移动distance米
        #####################

        # 定义不同无人机的飞行航线
        def flight_plan_1(node):
            t1.takeoff(node)
            t1.move_forward(node, 10.0) # 向前移动10米
            t1.move_left(node, 5.0) # 向左移动5米
            t1.move_forward(node, 10.0) # 向前移动10米
            t1.move_right(node, 5.0) # 向左移动5米

        def flight_plan_2(node):
            t1.takeoff(node)
            t1.move_left(node, 10.0)
            t1.move_forward(node, 15.0)
        
        def flight_plan_3(node):
            t1.takeoff(node)
            t1.move_forward(node, 10.0)
            t1.move_left(node, 25.0)
            
        def flight_plan_4(node):
            t1.takeoff(node)
            t1.move_forward(node, 15.0)
            t1.move_right(node, 5.0)

        def flight_plan_5(node):
            t1.takeoff(node)
            t1.move_forward(node, 30.0)
            t1.move_left(node, 15.0)

        # 飞行计划列表，每个元素对应一个无人机的飞行路线
        flight_plan_funcs = [
            flight_plan_1,
            flight_plan_2,
            flight_plan_3,
            flight_plan_4,
            flight_plan_5
        ]

        t1.execute_flight_plan(flight_plan_funcs)

        # 打印所有无人机的最终位置
        for node, index, position in t1.positions:
            print(f"[UAV {index}] Final position: x={position[0]:.2f}, y={position[1]:.2f}, z={position[2]:.2f}")

    except KeyboardInterrupt:
        print("\n--Simulation interrupted by user.")
    except Exception as e:
        print(f"\n--An error occurred: {e}")

import argparse
import time
from swarmae.SwarmAEClient import SwarmAEClient
from concurrent.futures import ThreadPoolExecutor
import math

class SwarmaeClass:
    def __init__(self):
        self.ae_client = None
        self.vehicles = []
        self.vehicle_info = {}  # 存储每辆车的初始位置和其他信息
        self.positions = []  # 存储每辆车的当前位置

    def create_client(self, ue_ip="127.0.0.1", ue_port=2000):
        print("--Connecting to UE4 simulation:", ue_ip, ue_port)
        self.ae_client = SwarmAEClient(ue_ip=ue_ip, ue_port=ue_port)
        return self.ae_client

    def create_vehicle(self, id, max_retries=3):
        for attempt in range(max_retries):
            frame_timestamp = int(round(time.time() * 1000))
            frame_timestamp, sw_vehicle, sw_code = self.ae_client.register_node(
                node_type="四轮车", node_name="无人车" + str(id), node_no=id, frame_timestamp=frame_timestamp
            )
            
            if sw_code == 200:
                vehicle_name, vehicle_no, team, vehicle_type, _, _ = sw_vehicle.get_node_info()
                initial_location = sw_vehicle.get_location()[:3]
                print(f"--Vehicle {id} registered successfully: {vehicle_name}, Type: {vehicle_type}")
                print(f"--Initial position: x={initial_location[0]:.2f}, y={initial_location[1]:.2f}, z={initial_location[2]:.2f}")
                self.vehicle_info[sw_vehicle] = {
                    'initial_location': initial_location
                }
                self.positions.append((sw_vehicle, id, initial_location))  # 保存车辆及其ID
                self.vehicles.append(sw_vehicle)
                return sw_code, sw_vehicle
            else:
                print(f"--Vehicle {id} registration failed with code {sw_code}, attempt {attempt+1} of {max_retries}")
                time.sleep(2)  # 等待2秒后重试

        print(f"Vehicle {id} registration failed after {max_retries} attempts.")
        return sw_code, None

    def move_vehicle(self, vehicle, throttle, steer=0.0):
        """控制车辆的移动和转向"""
        vehicle.apply_control(throttle, steer, 0.0, 0, 1)
        # vehicle.set_vehicle_steer(steer)  # 设置车辆方向盘转角
        # vehicle.control_vehicle(throttle=throttle)  # 控制车辆油门
        print(f"【Vehicle {self.vehicles.index(vehicle) + 1}】 is moving with throttle={throttle}, steer={steer}")

    def brake_vehicle(self, vehicle):
        """控制车辆刹车并确保车辆停止"""
        vehicle.apply_control(0.0, 0.0, 0.0, 1, 0)
        # while True:
        #     # 获取车辆当前速度信息
        #     velocity_x, velocity_y, velocity_z, velocity, _, _, _, _, _, p, q, r, _ = vehicle.get_velocity()
        #     # speed = math.sqrt(velocity_x**2 + velocity_y**2 + velocity_z**2)  # 计算总速度
        #     # print(velocity_x, velocity_y)
            
        #     if velocity_x <= 0.0 or velocity_y <= 0.0:  # 如果速度接近零，则停止车辆
        #         vehicle.apply_control(0.0, 0.0, 0.0, 1, 0)
        #         # vehicle.control_vehicle(throttle=0.0)
        #         print(f"【Vehicle {self.vehicles.index(vehicle) + 1}】 has stopped.")
        #         break
        #     else:
        #         vehicle.control_vehicle(throttle=-1.0)  # 否则继续减速
        #     time.sleep(0.1)  # 等待0.1秒再次检查速度

    def update_position(self, vehicle):
        """更新并保存车辆的当前位置。"""
        current_location = vehicle.get_location()[:3]
        vehicle_index = self.vehicles.index(vehicle)
        self.positions[vehicle_index] = (vehicle, vehicle_index + 1, current_location)

    def execute_vehicle_movement(self, movement_func, vehicle):
        """执行单个车辆的移动计划。"""
        if movement_func:
            movement_func(vehicle)  # 执行传入的移动计划函数
        else:
            print(f"No movement plan provided for vehicle {vehicle}. Skipping control.")

    def execute_movement_plans(self, movement_plan_funcs):
        """执行所有车辆的移动计划。"""
        if len(self.vehicles) < len(movement_plan_funcs):
            print("Not all vehicles were successfully registered. Adjusting the movement plans.")
            movement_plan_funcs = movement_plan_funcs[:len(self.vehicles)]

        with ThreadPoolExecutor() as executor:
            futures = []
            for vehicle, movement_plan_func in zip(self.vehicles, movement_plan_funcs):
                futures.append(executor.submit(self.execute_vehicle_movement, movement_plan_func, vehicle))
            for future in futures:
                future.result()  # 等待所有车辆完成其移动计划

if __name__ == '__main__':
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("-i", "--ip", default="127.0.0.1", help="指定服务器IP", type=str)
        parser.add_argument("-p", "--port", default=2000, help="指定服务器Port", type=int)
        parser.add_argument("-n", "--number", default=2, help="指定车辆数量", type=int)
        args = parser.parse_args()

        t1 = SwarmaeClass()
        t1.create_client(args.ip, args.port)
        for i in range(args.number):
            t1.create_vehicle(i+1)
        

        #####################
        # 定义移动计划
        # 前面的部分不需要改，要控制车辆的移动，只需修改下面的部分即可
        # movement_plan_n即第n个车辆的移动计划函数
        # t1.move_vehicle(vehicle, throttle, steer)  # 移动
        #####################

        def movement_plan_1(vehicle):
            t1.move_vehicle(vehicle, throttle=1.0)  # 向前移动
            time.sleep(5)  # 移动5秒
            t1.brake_vehicle(vehicle)  # 刹车并停止
            t1.move_vehicle(vehicle, throttle=0.0, steer=-0.5)  # 向左转
            time.sleep(5)
            t1.brake_vehicle(vehicle)  # 刹车并停止
            t1.move_vehicle(vehicle, throttle=1.0)
            time.sleep(4)
            t1.brake_vehicle(vehicle)  # 刹车并停止

        def movement_plan_2(vehicle):
            t1.move_vehicle(vehicle, throttle=0.1, steer=-1.0)  # 向前向左移动
            time.sleep(5)
            t1.move_vehicle(vehicle, throttle=1.0, steer=0)  # 向前向左移动
            time.sleep(5)
            t1.brake_vehicle(vehicle)  # 刹车并停止

        # 移动计划列表，每个元素对应一个车辆的移动路线
        movement_plan_funcs = [
            movement_plan_1,
            movement_plan_2,
            # 添加更多车辆的移动计划
        ]

        t1.execute_movement_plans(movement_plan_funcs)

        # 打印所有车辆的最终位置
        for vehicle, index, position in t1.positions:
            print(f"【Vehicle {index}】 Final position: x={position[0]:.2f}, y={position[1]:.2f}, z={position[2]:.2f}")

    except KeyboardInterrupt:
        print("\n--Simulation interrupted by user.")
    except Exception as e:
        print(f"\n--An error occurred: {e}")

#!/usr/bin/env python
import argparse
import time

map_info = {}

def get_subject_one_info(world):
    subject_one_info = {
        'Bridges': {
            'location': []
        },
        'SpeedBumps': {
            'location': []
        }
    }

    # 获取桥的坐标信息
    BP_Bridges = list(filter(lambda x: x.name.startswith('BP_Bridge'), world.get_environment_objects()))
    for bridge in BP_Bridges:
        if hasattr(bridge, 'transform') and hasattr(bridge.transform, 'location'):
            subject_one_info['Bridges']['location'].append([
                bridge.transform.location.x,
                bridge.transform.location.y,
                bridge.transform.location.z
            ])

    # 获取区域线的坐标信息
    SM_SpeedBumps = list(filter(lambda x: x.name.startswith('SM_SpeedBump'), world.get_environment_objects()))
    for speed_bump in SM_SpeedBumps:
        if hasattr(speed_bump, 'transform') and hasattr(speed_bump.transform, 'location'):
            subject_one_info['SpeedBumps']['location'].append([
                speed_bump.transform.location.x,
                speed_bump.transform.location.y,
                speed_bump.transform.location.z
            ])

    return subject_one_info

def get_subject_two_info(world):
    subject_two_info = {
        'FireCoverArea': {
            'location': []
        },
        'Obstacles': {
            'location': []
        }
    }

    # 获取火力掩护区域的坐标信息
    SM_Walls = list(filter(lambda x: x.name.startswith('SM_Wall_Single'), world.get_environment_objects()))
    for wall in SM_Walls:
        if hasattr(wall, 'transform') and hasattr(wall.transform, 'location'):
            subject_two_info['FireCoverArea']['location'].append([
                wall.transform.location.x,
                wall.transform.location.y,
                wall.transform.location.z
            ])

    # 获取路障的坐标信息
    SM_Obstacles = list(filter(lambda x: x.name.startswith('SM_Obstacle'), world.get_environment_objects()))
    for obstacle in SM_Obstacles:
        if hasattr(obstacle, 'transform') and hasattr(obstacle.transform, 'location'):
            subject_two_info['Obstacles']['location'].append([
                obstacle.transform.location.x,
                obstacle.transform.location.y,
                obstacle.transform.location.z
            ])

    return subject_two_info

def get_subject_four_info(world):
    subject_four_info = {
        'Waypoints': {
            'location': []
        }
    }

    # 获取航路点的坐标信息
    Waypoints = list(filter(lambda x: x.name.startswith('Waypoint'), world.get_environment_objects()))
    for waypoint in Waypoints:
        if hasattr(waypoint, 'transform') and hasattr(waypoint.transform, 'location'):
            subject_four_info['Waypoints']['location'].append([
                waypoint.transform.location.x,
                waypoint.transform.location.y,
                waypoint.transform.location.z
            ])

    return subject_four_info
            
def get_subject_five_info(world):
    subject_five_info = {
        'OuterDoorFrames': {
            'location': [],
            'rotation': []
        },
        'InnerDoorFrames': {
            'location': [],
            'rotation': []
        }
    }

    # 获取外门框的位置信息和角度信息
    SM_DoorFrames = list(filter(lambda x: x.name.startswith('SM_DoorFrame'), world.get_environment_objects()))
    for door_frame in SM_DoorFrames:
        if hasattr(door_frame, 'transform') and hasattr(door_frame.transform, 'location') and hasattr(door_frame.transform, 'rotation'):
            subject_five_info['OuterDoorFrames']['location'].append([
                door_frame.transform.location.x,
                door_frame.transform.location.y,
                door_frame.transform.location.z
            ])
            subject_five_info['OuterDoorFrames']['rotation'].append([
                door_frame.transform.rotation.pitch,
                door_frame.transform.rotation.yaw,
                door_frame.transform.rotation.roll
            ])

    # 获取内门框的位置信息和角度信息
    BP_DoorInsts = list(filter(lambda x: x.name.startswith('BP_DoorInst'), world.get_environment_objects()))
    for door_inst in BP_DoorInsts:
        if hasattr(door_inst, 'transform') and hasattr(door_inst.transform, 'location') and hasattr(door_inst.transform, 'rotation'):
            subject_five_info['InnerDoorFrames']['location'].append([
                door_inst.transform.location.x,
                door_inst.transform.location.y,
                door_inst.transform.location.z
            ])
            subject_five_info['InnerDoorFrames']['rotation'].append([
                door_inst.transform.rotation.pitch,
                door_inst.transform.rotation.yaw,
                door_inst.transform.rotation.roll
            ])

    return subject_five_info

def main():
    argparser = argparse.ArgumentParser(description=__doc__)
    argparser.add_argument('-i', '--address', default='127.0.0.1')
    argparser.add_argument('-p', '--port',type=int,default=2000)
    argparser.add_argument('-n', '--number',type=int,default=1)
    argparser.add_argument('-s', '--subject',type=int,default=1)
    args = argparser.parse_args()

    ue_ip = args.address.strip()
    ue_port = args.port
    num = args.number
    try:
        # 1. 首先引入sdk
        from swarmae.SwarmAEClient import SwarmAEClient
        from swarmae.SwarmAEWorld import SwarmAEWorld
        # 2. 连接UE
        client = SwarmAEClient(ue_ip=ue_ip, ue_port=ue_port)
        # sdk读取仿真世界，需要 1.0.1 sdk 才支持
        timestamp, world, code = client.get_world()
        
        # 获取科目的地图信息
        if args.subject == 1:
            map_info = get_subject_one_info(world)
        elif args.subject == 2:
            map_info = get_subject_two_info(world)
        elif args.subject == 4:
            map_info = get_subject_four_info(world)
        elif args.subject == 5:
            map_info = get_subject_five_info(world)
        
        print(map_info)

    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass