import os
import logging
import numpy as np
from typing import Dict
import pandas as pd


class MergeObj:
    def __init__(self, output_dir: str):
        self.output_dir = output_dir
        self.logger = logging.getLogger('UAV_Preprocess.MergeObj')

    def merge_two_objs(self, obj1_path: str, obj2_path: str, output_path: str):
        """使用Open3D合并两个OBJ文件"""
        try:
            self.logger.info("开始合并OBJ模型")
            self.logger.info(f"输入模型1: {obj1_path}")
            self.logger.info(f"输入模型2: {obj2_path}")
            self.logger.info(f"输出模型: {output_path}")

            # 检查输入文件是否存在
            if not os.path.exists(obj1_path) or not os.path.exists(obj2_path):
                raise FileNotFoundError("输入模型文件不存在")

            # 读取OBJ文件
            mesh1 = o3d.io.read_triangle_mesh(obj1_path)
            mesh2 = o3d.io.read_triangle_mesh(obj2_path)

            if mesh1.is_empty() or mesh2.is_empty():
                raise ValueError("无法读取OBJ文件或文件为空")

            # # 计算并对齐中心点
            # center1 = mesh1.get_center()
            # center2 = mesh2.get_center()
            # translation_vector = center2 - center1
            # mesh2.translate(translation_vector)

            # 不对齐，直接合并网格
            combined_mesh = mesh1 + mesh2

            # 优化合并后的网格
            combined_mesh.remove_duplicated_vertices()
            combined_mesh.remove_duplicated_triangles()
            combined_mesh.compute_vertex_normals()

            # 保存合并后的模型
            if not o3d.io.write_triangle_mesh(output_path, combined_mesh):
                raise RuntimeError("保存合并后的模型失败")

            self.logger.info(f"模型合并成功，已保存至: {output_path}")

        except Exception as e:
            self.logger.error(f"合并OBJ模型时发生错误: {str(e)}", exc_info=True)
            raise

    def merge_grid_obj(self, grid_points: Dict[int, pd.DataFrame]):
        """合并所有网格的OBJ模型"""
        self.logger.info("开始合并所有网格的OBJ模型")

        if len(grid_points) < 2:
            self.logger.info("只有一个网格，无需合并")
            return

        input_obj1, input_obj2 = None, None
        merge_count = 0

        try:
            for grid_idx, points in grid_points.items():
                grid_obj = os.path.join(
                    self.output_dir,
                    f"grid_{grid_idx + 1}",
                    "project",
                    "odm_texturing_25d",
                    "odm_textured_model_geo.obj"
                )

                if not os.path.exists(grid_obj):
                    self.logger.warning(
                        f"网格 {grid_idx + 1} 的OBJ文件不存在: {grid_obj}")
                    continue

                if input_obj1 is None:
                    input_obj1 = grid_obj
                    self.logger.info(f"设置第一个输入OBJ: {input_obj1}")
                else:
                    input_obj2 = grid_obj
                    output_obj = os.path.join(
                        self.output_dir, "merged_model.obj")

                    self.logger.info(
                        f"开始合并第 {merge_count + 1} 次:\n"
                        f"输入1: {input_obj1}\n"
                        f"输入2: {input_obj2}\n"
                        f"输出: {output_obj}"
                    )

                    self.merge_two_objs(input_obj1, input_obj2, output_obj)
                    merge_count += 1

                    input_obj1 = output_obj
                    input_obj2 = None

            self.logger.info(
                f"OBJ模型合并完成，共执行 {merge_count} 次合并，"
                f"最终输出文件: {input_obj1}"
            )

        except Exception as e:
            self.logger.error(f"OBJ模型合并过程中发生错误: {str(e)}", exc_info=True)
            raise


if __name__ == "__main__":
    import sys
    sys.path.append(os.path.dirname(
        os.path.dirname(os.path.abspath(__file__))))
    from utils.logger import setup_logger
    import pandas as pd

    # 设置输出目录和日志
    output_dir = r"G:\ODM_output\1009"
    setup_logger(output_dir)

    # 构造测试用的grid_points字典
    # 假设我们有两个网格，每个网格包含一些GPS点的DataFrame
    grid_points = {
        0: pd.DataFrame({
            'latitude': [39.9, 39.91],
            'longitude': [116.3, 116.31],
            'altitude': [100, 101]
        }),
        1: pd.DataFrame({
            'latitude': [39.92, 39.93],
            'longitude': [116.32, 116.33],
            'altitude': [102, 103]
        })
    }

    # 创建MergeObj实例并执行合并
    merge_obj = MergeObj(output_dir)
    merge_obj.merge_grid_obj(grid_points)