import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import matplotlib.pyplot as plt
from datetime import timedelta
import logging
import numpy as np
from utils.gps_extractor import GPSExtractor
from utils.logger import setup_logger

class GPSTimeVisualizer:
    """按时间组可视化GPS点"""
    
    def __init__(self, image_dir: str, output_dir: str):
        self.image_dir = image_dir
        self.output_dir = output_dir
        self.logger = logging.getLogger('UAV_Preprocess.GPSVisualizer')
        
    def _group_by_time(self, points_df, time_threshold=timedelta(minutes=5)):
        """按时间间隔对点进行分组"""
        if 'date' not in points_df.columns:
            self.logger.error("数据中缺少date列")
            return [points_df]

        # 将date为空的行单独作为一组
        null_date_group = points_df[points_df['date'].isna()]
        valid_date_points = points_df[points_df['date'].notna()]

        if not null_date_group.empty:
            self.logger.info(f"发现 {len(null_date_group)} 个无时间戳的点，将作为单独分组")

        if valid_date_points.empty:
            self.logger.warning("没有有效的时间戳数据")
            return [null_date_group] if not null_date_group.empty else []

        # 按时间排序
        valid_date_points = valid_date_points.sort_values('date')
        
        # 计算时间差
        time_diffs = valid_date_points['date'].diff()
        
        # 找到时间差超过阈值的位置
        time_groups = []
        current_group_start = 0
        
        for idx, time_diff in enumerate(time_diffs):
            if time_diff and time_diff > time_threshold:
                # 添加当前组
                current_group = valid_date_points.iloc[current_group_start:idx]
                time_groups.append(current_group)
                current_group_start = idx
        
        # 添加最后一组
        last_group = valid_date_points.iloc[current_group_start:]
        if not last_group.empty:
            time_groups.append(last_group)
            
        # 如果有空时间戳的点，将其作为最后一组
        if not null_date_group.empty:
            time_groups.append(null_date_group)
            
        return time_groups

    def visualize_time_groups(self, time_threshold=timedelta(minutes=5)):
        """在同一张图上显示所有时间组，用不同颜色区分"""
        # 提取GPS数据
        extractor = GPSExtractor(self.image_dir)
        gps_points = extractor.extract_all_gps()
        
        # 按时间分组
        time_groups = self._group_by_time(gps_points, time_threshold)
        
        # 创建图形
        plt.figure(figsize=(15, 10))
        
        # 生成不同的颜色
        colors = plt.cm.rainbow(np.linspace(0, 1, len(time_groups)))
        
        # 为每个时间组绘制点和轨迹
        for idx, (group, color) in enumerate(zip(time_groups, colors)):
            if not group['date'].isna().any():
                # 有时间戳的组
                sorted_group = group.sort_values('date')
                
                # 绘制轨迹线
                plt.plot(sorted_group['lon'], sorted_group['lat'],
                        color=color, linestyle='-', linewidth=1.5, alpha=0.6,
                        label=f'Flight Path {idx + 1}')
                
                # 绘制GPS点
                plt.scatter(sorted_group['lon'], sorted_group['lat'],
                          color=color, marker='o', s=30, alpha=0.6)
                
                # 标记起点和终点
                plt.scatter(sorted_group['lon'].iloc[0], sorted_group['lat'].iloc[0],
                          color=color, marker='^', s=100,
                          label=f'Start {idx + 1} ({sorted_group["date"].min().strftime("%H:%M:%S")})')
                plt.scatter(sorted_group['lon'].iloc[-1], sorted_group['lat'].iloc[-1],
                          color=color, marker='s', s=100,
                          label=f'End {idx + 1} ({sorted_group["date"].max().strftime("%H:%M:%S")})')
            else:
                # 无时间戳的组
                plt.scatter(group['lon'], group['lat'],
                          color=color, marker='x', s=50, alpha=0.6,
                          label='No Timestamp Points')
        
        plt.title("GPS Points by Time Groups", fontsize=14)
        plt.xlabel("Longitude", fontsize=12)
        plt.ylabel("Latitude", fontsize=12)
        plt.grid(True)
        
        # 调整图例位置和大小
        plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=10)
        
        # 调整布局以适应图例
        plt.tight_layout()
        
        # 保存图片
        plt.savefig(os.path.join(self.output_dir, 'gps_time_groups_combined.png'), 
                   dpi=300, bbox_inches='tight')
        plt.close()
        
        self.logger.info(f"已生成包含 {len(time_groups)} 个时间组的组合可视化图形")


if __name__ == '__main__':
    # 设置数据集路径
    DATASET = r'F:\error_data\20241108134711\3D'
    output_dir = r'E:\studio2\ODM_pro\test'
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置日志
    setup_logger(os.path.dirname(output_dir))
    
    # 创建可视化器并生成图形
    visualizer = GPSTimeVisualizer(DATASET, output_dir)
    visualizer.visualize_time_groups(time_threshold=timedelta(minutes=5))