#!/usr/bin/env python3
"""
飛行任務規劃模組
負責將 GPS 點擊座標轉換為隊形飛行任務
"""
import math
from typing import List, Tuple, Optional


class CoordinateConverter:
    """GPS 座標與 Local 座標的轉換器"""
    
    def __init__(self, origin_lat: float, origin_lon: float, origin_alt: float = 0):
        """
        初始化座標轉換器
        
        Args:
            origin_lat: 參考原點緯度
            origin_lon: 參考原點經度
            origin_alt: 參考原點高度（公尺，相對於海平面）
        """
        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]:
        """
        GPS 座標轉換為 Local 座標（ENU 系統：East-North-Up）
        
        Args:
            lat: 緯度
            lon: 經度
            alt: 高度（公尺，相對於海平面）
        
        Returns:
            (x, y, z): Local 座標（公尺）
                x: East（東方）
                y: North（北方）
                z: Up（向上）
        """
        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
        
        # 使用 Equirectangular projection（適用於小範圍 < 10 km）
        x = self.R * dlon * math.cos((lat_rad + origin_lat_rad) / 2)  # East
        y = self.R * dlat  # North
        z = alt - self.origin_alt  # Up
        
        return x, y, z
    
    def local_to_gps(self, x: float, y: float, z: float) -> Tuple[float, float, float]:
        """
        Local 座標轉換為 GPS 座標
        
        Args:
            x: East（東方，公尺）
            y: North（北方，公尺）
            z: Up（向上，公尺）
        
        Returns:
            (lat, lon, alt): GPS 座標
        """
        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):
        """
        初始化隊形規劃器
        
        Args:
            spacing: 無人機間距（公尺）
            base_altitude: 基準飛行高度（公尺，相對於參考原點）
            altitude_diff: M 字形的高低差（公尺）
        """
        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]) -> Tuple[List[Tuple[float, float, float]], 
                                                                                List[Tuple[float, float, float]]]:
        """
        規劃兩階段隊形任務
        
        Args:
            drone_gps_positions: 當前無人機 GPS 位置列表 [(lat, lon, alt), ...]
            target_gps: 目標點 GPS 座標 (lat, lon, alt)
        
        Returns:
            stage1_gps: 階段 1（集合點）的 GPS 航點列表
            stage2_gps: 階段 2（目標點）的 GPS 航點列表
            origin: 中心點（參考原點）的 GPS 座標 (lat, lon, alt)
            origin: 中心點（參考原點）的 GPS 座標 (lat, lon, alt)
        """
        if len(drone_gps_positions) == 0:
            raise ValueError("無人機位置列表不能為空")
        
        # ===== 步驟 1: 計算中央點並設為參考原點 =====
        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)
        
        print(f"📍 參考原點: ({center_lat:.6f}, {center_lon:.6f}, {center_alt:.1f}m)")
        
        # ===== 步驟 2: 轉換到 Local 座標系 =====
        drone_local_positions = []
        for lat, lon, alt in drone_gps_positions:
            x, y, z = self.converter.gps_to_local(lat, lon, alt)
            drone_local_positions.append((x, y, z))
        
        target_local = self.converter.gps_to_local(
            target_gps[0], target_gps[1], target_gps[2]
        )
        
        # ===== 步驟 3: 在 Local 座標系中計算隊形 =====
        stage1_local, stage2_local = self._calculate_formation_local(
            drone_local_positions,
            target_local
        )
        
        # ===== 步驟 4: 轉回 GPS 座標 =====
        stage1_gps = []
        for x, y, z in stage1_local:
            lat, lon, alt = self.converter.local_to_gps(x, y, z)
            stage1_gps.append((lat, lon, alt))
        
        stage2_gps = []
        for x, y, z in stage2_local:
            lat, lon, alt = self.converter.local_to_gps(x, y, z)
            stage2_gps.append((lat, lon, alt))
        
        print(f"✅ 任務規劃完成: {len(stage1_gps)} 台無人機，2 階段飛行")
        
        # ================================================================================
        # 【修改】回傳中心點座標供地圖視覺化使用
        # ================================================================================
        return stage1_gps, stage2_gps, self.current_origin
        # ================================================================================
    
    def _calculate_formation_local(self, 
                                   drone_positions: List[Tuple[float, float, float]], 
                                   target_point: Tuple[float, float, float]) -> Tuple[List[Tuple[float, float, float]], 
                                                                                       List[Tuple[float, float, float]]]:
        """
        在 Local 座標系中計算隊形
        
        Args:
            drone_positions: Local 座標的無人機位置 [(x, y, z), ...]
            target_point: Local 座標的目標點 (x, y, z)
        
        Returns:
            stage1_positions: 階段 1（中央點分佈）的 Local 座標
            stage2_positions: 階段 2（目標點分佈）的 Local 座標
        """
        N = len(drone_positions)
        
        # ===== 步驟 1: 計算中央點（在 Local 座標系中應該接近原點）=====
        C_x = sum(pos[0] for pos in drone_positions) / N
        C_y = sum(pos[1] for pos in drone_positions) / N
        C_z = sum(pos[2] for pos in drone_positions) / N
        
        print(f"   中央點 Local: ({C_x:.2f}, {C_y:.2f}, {C_z:.2f})")
        
        # ===== 步驟 2: 計算方向向量（中央點 → 目標點）=====
        T_x, T_y, T_z = target_point
        V_x = T_x - C_x
        V_y = T_y - C_y
        
        print(f"   方向向量: ({V_x:.2f}, {V_y:.2f})")
        
        # ===== 步驟 3: 計算垂直向量（逆時針旋轉 90°）=====
        P_x = -V_y
        P_y = V_x
        
        # 單位化垂直向量
        length = math.sqrt(P_x**2 + P_y**2)
        if length < 0.01:  # 避免除以零
            # 如果目標點與中央點重合，使用默認方向（向東）
            P_x_unit, P_y_unit = 1.0, 0.0
        else:
            P_x_unit = P_x / length
            P_y_unit = P_y / length
        
        print(f"   垂直單位向量: ({P_x_unit:.3f}, {P_y_unit:.3f})")
        
        # ===== 步驟 4: 計算無人機在垂直方向的投影並排序 =====
        projections = []
        for i, pos in enumerate(drone_positions):
            relative_x = pos[0] - C_x
            relative_y = pos[1] - C_y
            projection = relative_x * P_x_unit + relative_y * P_y_unit
            projections.append((projection, i))
        
        # 按投影值排序（保持左右順序）
        projections.sort()
        sorted_indices = [idx for _, idx in projections]
        
        print(f"   排序後的無人機索引: {sorted_indices}")
        
        # ===== 步驟 5: 分佈無人機位置 =====
        stage1_positions = [None] * N  # 預分配列表
        stage2_positions = [None] * N
        
        for rank, original_idx in enumerate(sorted_indices):
            # 計算相對中心的水平偏移
            offset = (rank - (N - 1) / 2) * self.spacing
            
            # 計算 M 字形高度（交替高低）
            if rank % 2 == 0:
                altitude = self.base_altitude + self.altitude_diff  # 高
            else:
                altitude = self.base_altitude - self.altitude_diff  # 低
            
            # 階段 1：在中央點附近分佈
            stage1_x = C_x + P_x_unit * offset
            stage1_y = C_y + P_y_unit * offset
            stage1_z = altitude
            
            # 階段 2：在目標點附近分佈（保持相同的相對位置）
            stage2_x = T_x + P_x_unit * offset
            stage2_y = T_y + P_y_unit * offset
            stage2_z = altitude
            
            # 按照原始索引存儲（保持無人機 ID 順序）
            stage1_positions[original_idx] = (stage1_x, stage1_y, stage1_z)
            stage2_positions[original_idx] = (stage2_x, stage2_y, stage2_z)
        
        return stage1_positions, stage2_positions
    
    def update_parameters(self, spacing: Optional[float] = None, 
                         base_altitude: Optional[float] = None, 
                         altitude_diff: Optional[float] = None):
        """
        更新隊形參數
        
        Args:
            spacing: 無人機間距（公尺）
            base_altitude: 基準飛行高度（公尺）
            altitude_diff: M 字形的高低差（公尺）
        """
        if spacing is not None:
            self.spacing = spacing
            print(f"   間距更新為: {spacing} m")
        
        if base_altitude is not None:
            self.base_altitude = base_altitude
            print(f"   基準高度更新為: {base_altitude} m")
        
        if altitude_diff is not None:
            self.altitude_diff = altitude_diff
            print(f"   高低差更新為: {altitude_diff} m")


# ===== 測試程式碼 =====
if __name__ == "__main__":
    print("=" * 60)
    print("隊形任務規劃器測試")
    print("=" * 60)
    
    # 模擬 5 台無人機的 GPS 位置
    drone_gps = [
        (24.123450, 120.567800, 100.0),  # 無人機 0
        (24.123470, 120.567820, 102.0),  # 無人機 1
        (24.123440, 120.567810, 98.0),   # 無人機 2
        (24.123460, 120.567830, 101.0),  # 無人機 3
        (24.123445, 120.567795, 99.0),   # 無人機 4
    ]
    
    # 目標點
    target_gps = (24.123600, 120.568000, 105.0)
    
    # 建立規劃器
    planner = FormationPlanner(
        spacing=5.0,           # 5 公尺間距
        base_altitude=10.0,    # 基準高度 10 公尺
        altitude_diff=2.0      # 高低差 2 公尺
    )
    
    # 規劃任務
    print("\n開始規劃任務...")
    stage1, stage2 = planner.plan_formation_mission(drone_gps, target_gps)
    
    # 顯示結果
    print("\n" + "=" * 60)
    print("階段 1：集合點位置（GPS）")
    print("=" * 60)
    for i, (lat, lon, alt) in enumerate(stage1):
        print(f"無人機 {i}: ({lat:.6f}, {lon:.6f}, {alt:.1f}m)")
    
    print("\n" + "=" * 60)
    print("階段 2：目標點位置（GPS）")
    print("=" * 60)
    for i, (lat, lon, alt) in enumerate(stage2):
        print(f"無人機 {i}: ({lat:.6f}, {lon:.6f}, {alt:.1f}m)")
    
    print("\n✅ 測試完成！")