import logging
from geopy.distance import geodesic
import matplotlib.pyplot as plt
import os


class GridDivider:
    """划分九宫格，并将图片分配到对应网格"""

    def __init__(self, overlap=0.1, output_dir=None):
        self.overlap = overlap
        self.output_dir = output_dir
        self.logger = logging.getLogger('UAV_Preprocess.GridDivider')
        self.logger.info(f"初始化网格划分器，重叠率: {overlap}")

    def divide_grids(self, points_df, grid_size=500):
        """计算边界框并划分九宫格"""
        self.logger.info("开始划分九宫格")

        min_lat, max_lat = points_df['lat'].min(), points_df['lat'].max()
        min_lon, max_lon = points_df['lon'].min(), points_df['lon'].max()

        # 计算区域的实际距离（米）
        width = geodesic((min_lat, min_lon), (min_lat, max_lon)).meters
        height = geodesic((min_lat, min_lon), (max_lat, min_lon)).meters

        self.logger.info(
            f"区域宽度: {width:.2f}米, 高度: {height:.2f}米"
        )

        # 计算需要划分的网格数量
        num_grids_width = int(
            width / grid_size) if int(width / grid_size) > 0 else 1
        num_grids_height = int(
            height / grid_size) if int(height / grid_size) > 0 else 1

        # 计算每个网格对应的经纬度步长
        lat_step = (max_lat - min_lat) / num_grids_height
        lon_step = (max_lon - min_lon) / num_grids_width

        grids = []
        for i in range(num_grids_height):
            for j in range(num_grids_width):
                grid_min_lat = min_lat + i * lat_step - self.overlap * lat_step
                grid_max_lat = min_lat + \
                    (i + 1) * lat_step + self.overlap * lat_step
                grid_min_lon = min_lon + j * lon_step - self.overlap * lon_step
                grid_max_lon = min_lon + \
                    (j + 1) * lon_step + self.overlap * lon_step
                grids.append((grid_min_lat, grid_max_lat,
                             grid_min_lon, grid_max_lon))

                self.logger.debug(
                    f"网格[{i},{j}]: 纬度[{grid_min_lat:.6f}, {grid_max_lat:.6f}], "
                    f"经度[{grid_min_lon:.6f}, {grid_max_lon:.6f}]"
                )

        self.logger.info(
            f"成功划分为 {len(grids)} 个网格 ({num_grids_width}x{num_grids_height})")

        # 添加可视化调用
        self.visualize_grids(points_df, grids)

        return grids

    def assign_to_grids(self, points_df, grids):
        """将点分配到对应网格"""
        self.logger.info(f"开始将 {len(points_df)} 个点分配到网格中")

        grid_points = {i: [] for i in range(len(grids))}
        points_assigned = 0
        multiple_grid_points = 0

        for _, point in points_df.iterrows():
            point_assigned = False
            for i, (min_lat, max_lat, min_lon, max_lon) in enumerate(grids):
                if min_lat <= point['lat'] <= max_lat and min_lon <= point['lon'] <= max_lon:
                    grid_points[i].append(point.to_dict())
                    if point_assigned:
                        multiple_grid_points += 1
                    else:
                        points_assigned += 1
                        point_assigned = True

            self.logger.debug(
                f"点 {point['file']} (纬度: {point['lat']:.6f}, 经度: {point['lon']:.6f}) "
                f"被分配到网格"
            )

        # 记录每个网格的点数
        for grid_idx, points in grid_points.items():
            self.logger.info(f"网格 {grid_idx} 包含 {len(points)} 个点")

        self.logger.info(
            f"点分配完成: 总点数 {len(points_df)}, "
            f"成功分配 {points_assigned} 个点, "
            f"{multiple_grid_points} 个点被分配到多个网格"
        )

        return grid_points

    def visualize_grids(self, points_df, grids):
        """可视化网格划分和GPS点的分布"""
        self.logger.info("开始可视化网格划分")

        plt.figure(figsize=(12, 8))

        # 绘制GPS点
        plt.scatter(points_df['lon'], points_df['lat'],
                    c='blue', s=10, alpha=0.6, label='GPS点')

        # 绘制网格
        for i, (min_lat, max_lat, min_lon, max_lon) in enumerate(grids):
            plt.plot([min_lon, max_lon, max_lon, min_lon, min_lon],
                     [min_lat, min_lat, max_lat, max_lat, min_lat],
                     'r-', alpha=0.5)
            # 在网格中心添加网格编号
            center_lon = (min_lon + max_lon) / 2
            center_lat = (min_lat + max_lat) / 2
            plt.text(center_lon, center_lat, str(i),
                     horizontalalignment='center', verticalalignment='center')

        plt.title('网格划分与GPS点分布图')
        plt.xlabel('经度')
        plt.ylabel('纬度')
        plt.legend()
        plt.grid(True)

        # 如果提供了输出目录，保存图像
        if self.output_dir:
            save_path = os.path.join(
                self.output_dir, 'filter_imgs_visual', 'grid_division.png')
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"网格划分可视化图已保存至: {save_path}")

        plt.close()