from osgeo import gdal
import logging
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


class MergeTif:
    def __init__(self, input_tif1, input_tif2, output_tif):
        self.input_tif1 = input_tif1
        self.input_tif2 = input_tif2
        self.output_tif = output_tif
        self.logger = logging.getLogger('UAV_Preprocess.MergeTif')

    def merge(self):
        """合并两张TIF影像"""
        try:
            self.logger.info("开始合并TIF影像")
            self.logger.info(f"输入影像1: {self.input_tif1}")
            self.logger.info(f"输入影像2: {self.input_tif2}")
            self.logger.info(f"输出影像: {self.output_tif}")

            # 检查输入文件是否存在
            if not os.path.exists(self.input_tif1) or not os.path.exists(self.input_tif2):
                error_msg = "输入影像文件不存在"
                self.logger.error(error_msg)
                raise FileNotFoundError(error_msg)

            # 打开影像，检查投影是否一致
            datasets = [gdal.Open(tif)
                        for tif in [self.input_tif1, self.input_tif2]]
            if None in datasets:
                error_msg = "无法打开输入影像文件"
                self.logger.error(error_msg)
                raise ValueError(error_msg)

            projections = [dataset.GetProjection() for dataset in datasets]
            self.logger.debug(f"影像1投影: {projections[0]}")
            self.logger.debug(f"影像2投影: {projections[1]}")

            # 检查投影是否一致
            if len(set(projections)) != 1:
                error_msg = "影像的投影不一致，请先进行重投影！"
                self.logger.error(error_msg)
                raise ValueError(error_msg)

            # 创建 GDAL Warp 选项
            warp_options = gdal.WarpOptions(
                format="GTiff",
                resampleAlg="average",  # 设置重采样方法为平均值
                srcNodata=0,            # 输入影像中的无效值
                dstNodata=0,            # 输出影像中的无效值
                multithread=True        # 启用多线程优化
            )

            self.logger.info("开始执行影像拼接...")

            # 使用 GDAL 的 Warp 方法进行拼接
            result = gdal.Warp(
                self.output_tif,
                [self.input_tif1, self.input_tif2],  # 输入多张影像
                options=warp_options
            )

            if result is None:
                error_msg = "影像拼接失败"
                self.logger.error(error_msg)
                raise RuntimeError(error_msg)

            # 获取输出影像的基本信息
            output_dataset = gdal.Open(self.output_tif)
            if output_dataset:
                width = output_dataset.RasterXSize
                height = output_dataset.RasterYSize
                bands = output_dataset.RasterCount
                self.logger.info(f"拼接完成，输出影像大小: {width}x{height}，波段数: {bands}")

            self.logger.info(f"影像拼接成功，输出文件保存至: {self.output_tif}")

        except Exception as e:
            self.logger.error(f"影像拼接过程中发生错误: {str(e)}", exc_info=True)
            raise


if __name__ == "__main__":
    from utils.logger import setup_logger

    # 定义影像路径
    input_tif1 = r"G:\20240930091614\output\grid_1\project\odm_orthophoto\odm_orthophoto.tif"
    input_tif2 = r"G:\20240930091614\output\grid_2\project\odm_orthophoto\odm_orthophoto.tif"
    output_tif = r"G:\20240930091614\output\merged_orthophoto.tif"

    # 设置日志
    output_dir = r"E:\studio2\ODM_pro\test"
    setup_logger(output_dir)

    # 执行拼接
    merge_tif = MergeTif(input_tif1, input_tif2, output_tif)
    merge_tif.merge()