AirTrapMine/src/GUI/mission_planner.py

#!/usr/bin/env python3
"""
飛行任務規劃模組
支援: M-Formation, Circle, Leader-Follower (Bezier 轉彎), Grid Sweep
"""
import math
from typing import List, Tuple, Optional, Dict, Any
from enum import Enum


class MissionType(Enum):
    """任務類型"""
    M_FORMATION = "m_formation"
    CIRCLE_FORMATION = "circle_formation"
    LEADER_FOLLOWER = "leader_follower"
    GRID_SWEEP = "grid_sweep"


class CoordinateConverter:
    """GPS 座標與 Local 座標的轉換器"""

    def __init__(self, origin_lat: float, origin_lon: float, origin_alt: float = 0):
        self.origin_lat = origin_lat
        self.origin_lon = origin_lon
        self.origin_alt = origin_alt
        self.R = 6371000.0

    def gps_to_local(self, lat: float, lon: float, alt: float) -> Tuple[float, float, float]:
        lat_rad = math.radians(lat)
        lon_rad = math.radians(lon)
        origin_lat_rad = math.radians(self.origin_lat)
        origin_lon_rad = math.radians(self.origin_lon)

        dlat = lat_rad - origin_lat_rad
        dlon = lon_rad - origin_lon_rad

        x = self.R * dlon * math.cos((lat_rad + origin_lat_rad) / 2)
        y = self.R * dlat
        z = alt - self.origin_alt

        return x, y, z

    def local_to_gps(self, x: float, y: float, z: float) -> Tuple[float, float, float]:
        origin_lat_rad = math.radians(self.origin_lat)
        origin_lon_rad = math.radians(self.origin_lon)

        lat_rad = origin_lat_rad + (y / self.R)
        lon_rad = origin_lon_rad + (x / (self.R * math.cos((lat_rad + origin_lat_rad) / 2)))

        lat = math.degrees(lat_rad)
        lon = math.degrees(lon_rad)
        alt = self.origin_alt + z

        return lat, lon, alt


class FormationPlanner:
    """隊形規劃器"""

    def __init__(self, spacing: float = 5.0,
                 base_altitude: float = 10.0,
                 altitude_diff: float = 2.0):
        self.spacing = spacing
        self.base_altitude = base_altitude
        self.altitude_diff = altitude_diff
        self.current_origin = None
        self.converter = None

    def plan_formation_mission(self,
                               drone_gps_positions: List[Tuple[float, float, float]],
                               target_gps: Tuple[float, float, float],
                               mission_type: MissionType = MissionType.M_FORMATION,
                               params: Optional[Dict[str, Any]] = None
                               ) -> Tuple[List[List[Tuple[float, float, float]]],
                                          Tuple[float, float, float]]:
        if len(drone_gps_positions) == 0:
            raise ValueError("無人機位置列表不能為空")

        center_lat = sum(pos[0] for pos in drone_gps_positions) / len(drone_gps_positions)
        center_lon = sum(pos[1] for pos in drone_gps_positions) / len(drone_gps_positions)
        center_alt = sum(pos[2] for pos in drone_gps_positions) / len(drone_gps_positions)

        self.current_origin = (center_lat, center_lon, center_alt)
        self.converter = CoordinateConverter(center_lat, center_lon, center_alt)

        drone_local = [self.converter.gps_to_local(*pos) for pos in drone_gps_positions]
        target_local = self.converter.gps_to_local(*target_gps)

        if mission_type == MissionType.M_FORMATION:
            s1, s2 = self._calculate_m_formation(drone_local, target_local, params)
            waypoints_local = [[s1[i], s2[i]] for i in range(len(drone_local))]

        elif mission_type == MissionType.CIRCLE_FORMATION:
            s1, s2 = self._calculate_circle_formation(drone_local, target_local, params)
            waypoints_local = [[s1[i], s2[i]] for i in range(len(drone_local))]

        elif mission_type == MissionType.LEADER_FOLLOWER:
            params = params or {}
            route_wps_gps = params.get('route_waypoints')
            if route_wps_gps is None or len(route_wps_gps) < 2:
                raise ValueError("LEADER_FOLLOWER 需要至少 2 個路徑點")
            route_wps_local = [
                self.converter.gps_to_local(wp[0], wp[1], 0)[:2]
                for wp in route_wps_gps
            ]
            waypoints_local = self._calculate_leader_follower(drone_local, route_wps_local, params)

        elif mission_type == MissionType.GRID_SWEEP:
            params = params or {}
            rect_corners_gps = params.get('rect_corners')
            if rect_corners_gps is None or len(rect_corners_gps) != 4:
                raise ValueError("GRID_SWEEP 需要 4 個 GPS 角點")
            rect_corners_local = [
                self.converter.gps_to_local(c[0], c[1], 0)[:2]
                for c in rect_corners_gps
            ]
            waypoints_local = self._calculate_grid_sweep(drone_local, rect_corners_local, params)
        else:
            raise ValueError(f"不支援的任務類型: {mission_type}")

        waypoints_gps = []
        for drone_wps in waypoints_local:
            gps_wps = [self.converter.local_to_gps(*wp) for wp in drone_wps]
            waypoints_gps.append(gps_wps)

        return waypoints_gps, self.current_origin

    # ------------------------------------------------------------------ M-Formation

    def _calculate_m_formation(self, drone_positions, target_point, params):
        params = params or {}
        N = len(drone_positions)
        spacing = params.get('spacing', self.spacing)
        base_altitude = params.get('base_altitude', self.base_altitude)
        altitude_diff = params.get('altitude_diff', self.altitude_diff)

        C_x = sum(pos[0] for pos in drone_positions) / N
        C_y = sum(pos[1] for pos in drone_positions) / N

        T_x, T_y, T_z = target_point
        V_x, V_y = T_x - C_x, T_y - C_y
        P_x, P_y = -V_y, V_x

        length = math.sqrt(P_x ** 2 + P_y ** 2)
        P_x_unit, P_y_unit = (P_x / length, P_y / length) if length > 0.01 else (1.0, 0.0)

        projections = [((pos[0] - C_x) * P_x_unit + (pos[1] - C_y) * P_y_unit, i)
                       for i, pos in enumerate(drone_positions)]
        projections.sort()

        stage1_positions = [None] * N
        stage2_positions = [None] * N

        for rank, (_, original_idx) in enumerate(projections):
            offset = (rank - (N - 1) / 2) * spacing
            altitude = base_altitude + (altitude_diff if rank % 2 == 0 else -altitude_diff)
            stage1_positions[original_idx] = (C_x + P_x_unit * offset, C_y + P_y_unit * offset, altitude)
            stage2_positions[original_idx] = (T_x + P_x_unit * offset, T_y + P_y_unit * offset, altitude)

        return stage1_positions, stage2_positions

    # ------------------------------------------------------------------ Circle

    def _calculate_circle_formation(self, drone_positions, target_point, params):
        params = params or {}
        N = len(drone_positions)
        radius = params.get('radius', 10.0)
        altitude = params.get('altitude', 10.0)
        start_angle = params.get('start_angle', 0.0)

        center_x, center_y, center_z = target_point
        stage1_positions = []
        stage2_positions = []
        angle_step = 360.0 / N

        for i in range(N):
            angle_rad = math.radians(start_angle + angle_step * i)
            final_x = center_x + radius * math.cos(angle_rad)
            final_y = center_y + radius * math.sin(angle_rad)
            final_z = altitude

            current_x, current_y, current_z = drone_positions[i]
            stage1_positions.append((
                current_x + (final_x - current_x) * 0.5,
                current_y + (final_y - current_y) * 0.5,
                current_z + (final_z - current_z) * 0.5
            ))
            stage2_positions.append((final_x, final_y, final_z))

        return stage1_positions, stage2_positions

    # ------------------------------------------------------------------ 路徑跟隨 (Bezier 轉彎)

    def _calculate_leader_follower(self, drone_positions, route_wps_local, params):
        """
        路徑跟隨編隊 — Bezier 曲線轉彎版

        步驟:
            1. _build_center_path: 在轉折 WP 處用二次 Bezier 曲線平滑轉彎
            2. 固定排序，每架無人機沿中心路徑套用橫向+縱向偏移

        隊形 (俯視):
            |     (前進方向)         |
            |           D1          |  ← 左偏移, 後 0m
            |                D2     |  ← 右偏移, 後 5m
            |           D3          |  ← 左偏移, 後 10m
            |                D4     |  ← 右偏移, 後 15m
        """
        N = len(drone_positions)
        lateral_offset = params.get('lateral_offset', 3.0)
        longitudinal_spacing = params.get('longitudinal_spacing', 5.0)
        altitude = params.get('altitude', self.base_altitude)
        turn_margin = params.get('turn_margin', 0.35)      # 轉彎切入距離佔段長比例
        curve_resolution = params.get('curve_resolution', 8)  # 每個彎道的插值點數

        # Step 1: 建立帶 Bezier 轉彎的中心路徑
        center_path = self._build_center_path(
            route_wps_local, turn_margin, curve_resolution
        )

        # Step 2: 固定排序
        first_dir = (center_path[0][2], center_path[0][3])
        first_perp = (-first_dir[1], first_dir[0])
        C_x = sum(p[0] for p in drone_positions) / N
        C_y = sum(p[1] for p in drone_positions) / N

        projections = [
            ((pos[0] - C_x) * first_perp[0] + (pos[1] - C_y) * first_perp[1], i)
            for i, pos in enumerate(drone_positions)
        ]
        projections.sort()

        # Step 3: 偏移
        all_waypoints = [None] * N

        for rank, (_, original_idx) in enumerate(projections):
            lat_sign = -1 if rank % 2 == 0 else 1
            lat = lat_sign * lateral_offset
            lon = rank * longitudinal_spacing

            waypoints = []
            for (cx, cy, dx, dy) in center_path:
                perp_x, perp_y = -dy, dx
                ox = cx + lat * perp_x - lon * dx
                oy = cy + lat * perp_y - lon * dy
                waypoints.append((ox, oy, altitude))

            all_waypoints[original_idx] = waypoints

        return all_waypoints

    def _build_center_path(self, waypoints, turn_margin, curve_resolution):
        """
        建立帶 Bezier 曲線轉彎的中心路徑

        在每個轉折 WP 處:
            1. 計算 pre_turn = WP - d_in × T
            2. 計算 post_turn = WP + d_out × T
            3. 用二次 Bezier 曲線: B(t) = (1-t)²·pre + 2t(1-t)·WP + t²·post
            4. 方向從導數得到: B'(t) = 2(1-t)(WP-pre) + 2t(post-WP)

        Args:
            waypoints: 路徑航點 [(x, y), ...]
            turn_margin: 轉彎切入距離佔相鄰段長的比例 (0~0.5)
            curve_resolution: 每個彎道的 Bezier 插值點數

        Returns:
            [(x, y, dir_x, dir_y), ...] 中心路徑點 + 單位方向
        """
        num_wps = len(waypoints)

        if num_wps < 2:
            return [(waypoints[0][0], waypoints[0][1], 0.0, 1.0)]

        # 計算每段方向和長度
        segments = []
        for i in range(num_wps - 1):
            dx = waypoints[i + 1][0] - waypoints[i][0]
            dy = waypoints[i + 1][1] - waypoints[i][1]
            length = math.sqrt(dx ** 2 + dy ** 2)
            if length < 0.01:
                segments.append((0.0, 1.0, length))
            else:
                segments.append((dx / length, dy / length, length))

        path = []

        # 第一個航點
        path.append((waypoints[0][0], waypoints[0][1],
                      segments[0][0], segments[0][1]))

        # 中間航點 (轉彎處)
        for i in range(1, num_wps - 1):
            d_in_x, d_in_y, len_in = segments[i - 1]
            d_out_x, d_out_y, len_out = segments[i]

            # 切入距離 T: 取相鄰兩段中較短的 × turn_margin
            T = min(len_in, len_out) * turn_margin

            if T < 0.5:
                # 段太短，不插彎，直接加一個平均方向點
                avg_dx = d_in_x + d_out_x
                avg_dy = d_in_y + d_out_y
                avg_len = math.sqrt(avg_dx ** 2 + avg_dy ** 2)
                if avg_len > 0.01:
                    avg_dx /= avg_len
                    avg_dy /= avg_len
                else:
                    avg_dx, avg_dy = d_in_x, d_in_y
                path.append((waypoints[i][0], waypoints[i][1], avg_dx, avg_dy))
                continue

            # P0 = pre_turn (弧線起始)
            p0_x = waypoints[i][0] - d_in_x * T
            p0_y = waypoints[i][1] - d_in_y * T

            # P1 = WP 本身 (控制點)
            p1_x = waypoints[i][0]
            p1_y = waypoints[i][1]

            # P2 = post_turn (弧線結束)
            p2_x = waypoints[i][0] + d_out_x * T
            p2_y = waypoints[i][1] + d_out_y * T

            # 加入 pre_turn 點 (方向 = incoming)
            path.append((p0_x, p0_y, d_in_x, d_in_y))

            # 加入 Bezier 插值點
            for j in range(1, curve_resolution):
                t = j / curve_resolution

                # B(t) = (1-t)²·P0 + 2t(1-t)·P1 + t²·P2
                one_minus_t = 1.0 - t
                bx = one_minus_t * one_minus_t * p0_x + \
                     2 * t * one_minus_t * p1_x + \
                     t * t * p2_x
                by = one_minus_t * one_minus_t * p0_y + \
                     2 * t * one_minus_t * p1_y + \
                     t * t * p2_y

                # B'(t) = 2(1-t)(P1-P0) + 2t(P2-P1) → 切線方向
                tdx = 2 * one_minus_t * (p1_x - p0_x) + 2 * t * (p2_x - p1_x)
                tdy = 2 * one_minus_t * (p1_y - p0_y) + 2 * t * (p2_y - p1_y)

                # 正規化
                t_len = math.sqrt(tdx ** 2 + tdy ** 2)
                if t_len > 0.01:
                    tdx /= t_len
                    tdy /= t_len
                else:
                    tdx, tdy = d_in_x, d_in_y

                path.append((bx, by, tdx, tdy))

            # 加入 post_turn 點 (方向 = outgoing)
            path.append((p2_x, p2_y, d_out_x, d_out_y))

        # 最後一個航點
        path.append((waypoints[-1][0], waypoints[-1][1],
                      segments[-1][0], segments[-1][1]))

        return path

    # ------------------------------------------------------------------ 柵狀掃描

    def _calculate_grid_sweep(self, drone_positions, rect_corners_local, params):
        """柵狀掃描：掃描方向沿矩形長邊，切割方向沿短邊"""
        N = len(drone_positions)
        line_spacing = params.get('line_spacing', 5.0)
        altitude = params.get('altitude', self.base_altitude)

        c0, c1, c2, c3 = rect_corners_local

        edge_a = (c1[0] - c0[0], c1[1] - c0[1])
        len_a = math.sqrt(edge_a[0] ** 2 + edge_a[1] ** 2)
        edge_b = (c3[0] - c0[0], c3[1] - c0[1])
        len_b = math.sqrt(edge_b[0] ** 2 + edge_b[1] ** 2)

        if len_a >= len_b:
            sweep_vec = edge_a
            sweep_len = len_a
            cut_vec = edge_b
            cut_len = len_b
            sweep_start_mid = ((c0[0] + c3[0]) / 2, (c0[1] + c3[1]) / 2)
            sweep_end_mid = ((c1[0] + c2[0]) / 2, (c1[1] + c2[1]) / 2)
            cut_start_corner = c0
        else:
            sweep_vec = edge_b
            sweep_len = len_b
            cut_vec = edge_a
            cut_len = len_a
            sweep_start_mid = ((c0[0] + c1[0]) / 2, (c0[1] + c1[1]) / 2)
            sweep_end_mid = ((c3[0] + c2[0]) / 2, (c3[1] + c2[1]) / 2)
            cut_start_corner = c0

        sweep_dir = (sweep_vec[0] / sweep_len, sweep_vec[1] / sweep_len)
        cut_dir = (cut_vec[0] / cut_len, cut_vec[1] / cut_len)

        C_x = sum(p[0] for p in drone_positions) / N
        C_y = sum(p[1] for p in drone_positions) / N

        dist_to_start = math.sqrt(
            (C_x - sweep_start_mid[0]) ** 2 + (C_y - sweep_start_mid[1]) ** 2
        )
        dist_to_end = math.sqrt(
            (C_x - sweep_end_mid[0]) ** 2 + (C_y - sweep_end_mid[1]) ** 2
        )

        if dist_to_start <= dist_to_end:
            near_corner_a = cut_start_corner
        else:
            sweep_dir = (-sweep_dir[0], -sweep_dir[1])
            near_corner_a = (cut_start_corner[0] + sweep_vec[0],
                            cut_start_corner[1] + sweep_vec[1])

        projections = [
            ((pos[0] - C_x) * cut_dir[0] + (pos[1] - C_y) * cut_dir[1], i)
            for i, pos in enumerate(drone_positions)
        ]
        projections.sort()

        strip_width = cut_len / N
        drone_sweep_proj = C_x * sweep_dir[0] + C_y * sweep_dir[1]
        near_sweep_proj = near_corner_a[0] * sweep_dir[0] + near_corner_a[1] * sweep_dir[1]
        gather_sweep_proj = (drone_sweep_proj + near_sweep_proj) / 2

        all_waypoints = [None] * N

        for rank, (_, original_idx) in enumerate(projections):
            strip_center_offset = (rank + 0.5) * strip_width
            base_x = near_corner_a[0] + cut_dir[0] * strip_center_offset
            base_y = near_corner_a[1] + cut_dir[1] * strip_center_offset

            waypoints_list = []

            gather_offset = gather_sweep_proj - near_sweep_proj
            gx = base_x + sweep_dir[0] * gather_offset
            gy = base_y + sweep_dir[1] * gather_offset
            waypoints_list.append((gx, gy, altitude))

            num_lines = max(1, int(strip_width / line_spacing))
            actual_spacing = strip_width / num_lines
            first_line_offset = (rank * strip_width) + actual_spacing / 2

            entry_x = near_corner_a[0] + cut_dir[0] * first_line_offset
            entry_y = near_corner_a[1] + cut_dir[1] * first_line_offset
            waypoints_list.append((entry_x, entry_y, altitude))

            current_cut_offset = first_line_offset

            for line_idx in range(num_lines):
                line_near_x = near_corner_a[0] + cut_dir[0] * current_cut_offset
                line_near_y = near_corner_a[1] + cut_dir[1] * current_cut_offset
                line_far_x = line_near_x + sweep_dir[0] * sweep_len
                line_far_y = line_near_y + sweep_dir[1] * sweep_len

                if line_idx % 2 == 0:
                    waypoints_list.append((line_far_x, line_far_y, altitude))
                else:
                    waypoints_list.append((line_near_x, line_near_y, altitude))

                if line_idx < num_lines - 1:
                    next_cut_offset = current_cut_offset + actual_spacing
                    next_near_x = near_corner_a[0] + cut_dir[0] * next_cut_offset
                    next_near_y = near_corner_a[1] + cut_dir[1] * next_cut_offset
                    next_far_x = next_near_x + sweep_dir[0] * sweep_len
                    next_far_y = next_near_y + sweep_dir[1] * sweep_len

                    if line_idx % 2 == 0:
                        waypoints_list.append((next_far_x, next_far_y, altitude))
                    else:
                        waypoints_list.append((next_near_x, next_near_y, altitude))
                    current_cut_offset = next_cut_offset

            all_waypoints[original_idx] = waypoints_list

        return all_waypoints


# ================================================================================
# 測試
# ================================================================================
if __name__ == "__main__":
    planner = FormationPlanner()

    drones = [
        (24.123450, 120.567800, 0.0),
        (24.123470, 120.567820, 0.0),
        (24.123440, 120.567810, 0.0),
        (24.123460, 120.567830, 0.0),
    ]
    target = (24.12400, 120.56795, 10.0)

    # M-Formation
    wps, origin = planner.plan_formation_mission(drones, target, MissionType.M_FORMATION)
    print("M-Formation:")
    for i, wp_list in enumerate(wps):
        print(f"  Drone {i}: {len(wp_list)} waypoints")

    # Leader-Follower (Bezier 轉彎)
    route = [
        (24.12345, 120.56780),
        (24.12370, 120.56800),
        (24.12390, 120.56820),
        (24.12400, 120.56800),
        (24.12420, 120.56790),
    ]
    wps_lf, origin_lf = planner.plan_formation_mission(
        drones, target, MissionType.LEADER_FOLLOWER,
        params={
            'route_waypoints': route,
            'lateral_offset': 3.0,
            'longitudinal_spacing': 5.0,
            'turn_margin': 0.35,
            'curve_resolution': 8,
            'altitude': 10.0
        }
    )
    print(f"\nLeader-Follower (Bezier turns):")
    for i, wp_list in enumerate(wps_lf):
        print(f"  Drone {i}: {len(wp_list)} waypoints")
        for j, wp in enumerate(wp_list):
            print(f"    WP{j}: ({wp[0]:.6f}, {wp[1]:.6f}, {wp[2]:.1f})")

    # Grid Sweep
    rect = [
        (24.1237, 120.5679),
        (24.1237, 120.5683),
        (24.1240, 120.5683),
        (24.1240, 120.5679)
    ]
    wps2, origin2 = planner.plan_formation_mission(
        drones, target, MissionType.GRID_SWEEP,
        params={'rect_corners': rect, 'line_spacing': 5.0, 'altitude': 10.0}
    )
    print(f"\nGrid Sweep:")
    for i, wp_list in enumerate(wps2):
        print(f"  Drone {i}: {len(wp_list)} waypoints")