import os
import sys
import shutil
from pathlib import Path
import matplotlib.pyplot as plt
from matplotlib.widgets import RectangleSelector
import pandas as pd

# 添加项目根目录到系统路径
project_root = str(Path(__file__).parent.parent)
sys.path.append(project_root)

from utils.gps_extractor import GPSExtractor

class GPSSelector:
    def __init__(self, image_dir: str, output_dir: str = None):
        self.image_dir = image_dir
        self.output_dir = output_dir
        self.gps_points = None
        self.selected_points = []
        self.fig, self.ax = plt.subplots(figsize=(12, 8))
        self.scatter = None
        self.rs = None
        self.setup_plot()

    def extract_gps(self):
        """提取GPS数据"""
        extractor = GPSExtractor(self.image_dir)
        self.gps_points = extractor.extract_all_gps()
        print(f"成功提取 {len(self.gps_points)} 个GPS点")

    def setup_plot(self):
        """设置绘图"""
        self.ax.set_title('GPS Points - 使用鼠标拖动选择要删除的点')
        self.ax.set_xlabel('Longitude')
        self.ax.set_ylabel('Latitude')
        self.ax.grid(True)

        # 设置矩形选择器
        self.rs = RectangleSelector(
            self.ax, self.on_select,
            interactive=True,
            useblit=True,
            button=[1],  # 只响应左键
            props=dict(facecolor='red', alpha=0.3)
        )

        # 添加按钮回调
        self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)

    def plot_gps_points(self):
        """绘制GPS点"""
        if self.scatter is not None:
            self.scatter.remove()
        
        self.scatter = self.ax.scatter(
            self.gps_points['lon'],
            self.gps_points['lat'],
            c='blue',
            s=20,
            alpha=0.6
        )
        self.fig.canvas.draw_idle()

    def on_select(self, eclick, erelease):
        """矩形选择回调"""
        x1, y1 = eclick.xdata, eclick.ydata
        x2, y2 = erelease.xdata, erelease.ydata
        
        # 获取选中区域内的点
        mask = (
            (self.gps_points['lon'] >= min(x1, x2)) &
            (self.gps_points['lon'] <= max(x1, x2)) &
            (self.gps_points['lat'] >= min(y1, y2)) &
            (self.gps_points['lat'] <= max(y1, y2))
        )
        
        selected = self.gps_points[mask]
        self.selected_points.extend(selected['file'].tolist())
        
        # 从数据中移除选中的点
        self.gps_points = self.gps_points[~mask]
        
        # 更新绘图
        self.plot_gps_points()
        print(f"选中 {len(selected)} 个点，剩余 {len(self.gps_points)} 个点")

    def on_key_press(self, event):
        """键盘事件回调"""
        if event.key == 'enter':
            self.save_results()
            plt.close()
        elif event.key == 'escape':
            plt.close()

    def save_results(self):
        """保存结果"""
        if not self.output_dir:
            return

        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)
        removed_dir = os.path.join(self.output_dir, "removed_images")
        os.makedirs(removed_dir, exist_ok=True)

        # 移动被删除的图像
        for img_name in self.selected_points:
            src = os.path.join(self.image_dir, img_name)
            dst = os.path.join(removed_dir, img_name)
            shutil.move(src, dst)

        # 保存剩余点的信息
        self.gps_points.to_csv(
            os.path.join(self.output_dir, "remaining_points.csv"),
            index=False
        )
        
        print(f"已移动 {len(self.selected_points)} 张图片到 {removed_dir}")
        print(f"保留 {len(self.gps_points)} 个点")

    def run(self):
        """运行选择器"""
        self.extract_gps()
        self.plot_gps_points()
        plt.show()


if __name__ == "__main__":
    # 使用示例
    selector = GPSSelector(
        image_dir=r"E:\datasets\UAV\1619\project\images",
        output_dir=r"E:\datasets\UAV\1619\filtered"
    )
    selector.run()