# -*- coding: utf-8 -*-

__author__ = 'wanger'
__description__ = '按照掩膜范围裁剪栅格数据'
__date__ = '2024-11-25'
__copyright__ = '(C) 2024 by siwei'
__revision__ = '1.0'

import subprocess

from qgis._core import QgsProcessingParameterFolderDestination
from qgis.core import (
    QgsRasterLayer,
)

from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (QgsVectorLayer,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFile)


class ClipRasterByMaskProcessingAlgorithm(QgsProcessingAlgorithm):
    INPUT_SHP = 'INPUT_SHP'
    INPUT_RASTER = 'INPUT_RASTER'
    OUTPUT = 'OUTPUT'

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def createInstance(self):
        return ClipRasterByMaskProcessingAlgorithm()

    def name(self):
        return 'cliprasterbymask'

    def displayName(self):
        return self.tr('栅格裁剪')

    def group(self):
        return self.tr('栅格裁剪')

    def groupId(self):
        return 'rasterclip'

    def shortHelpString(self):
        return self.tr("遍历输入shp数据要素裁剪输入栅格，保存到输出文件夹。")

    def initAlgorithm(self, config=None):
        self.addParameter(QgsProcessingParameterFile(
            self.INPUT_SHP,
            '网格数据',
            extension='shp'
        ))
        self.addParameter(QgsProcessingParameterFile(
            self.INPUT_RASTER,
            '栅格数据',
            extension='tif'
        ))
        self.addParameter(QgsProcessingParameterFolderDestination(self.OUTPUT, '输出文件夹'))

    # 执行
    def processAlgorithm(self, parameters, context, feedback):
        input_shp = self.parameterAsString(parameters, self.INPUT_SHP, context)
        input_raster = self.parameterAsOutputLayer(parameters, self.INPUT_RASTER, context)
        output_raster = self.parameterAsString(parameters, self.OUTPUT, context)
        # 加载输入Shapefile和栅格数据
        vector_layer = QgsVectorLayer(input_shp, "vector_layer", "ogr")
        raster_layer = QgsRasterLayer(input_raster, "raster_layer")
        fields = vector_layer.fields()
        # 获取分辨率（像素大小）
        pixel_width = raster_layer.rasterUnitsPerPixelX()
        pixel_height = raster_layer.rasterUnitsPerPixelY()
        # 遍历输入图层的所有要素，并为每个要素保存一个单独的Shapefile文件
        total = 0
        for idx, feature in enumerate(vector_layer.getFeatures()):
            # 获取当前要素的几何数据
            geometry = feature.geometry()
            key = feature["新图号"]
            extent = geometry.boundingBox()
            # 计算裁剪范围的整数倍，确保裁剪范围对齐
            minX = extent.xMinimum()
            maxX = extent.xMaximum()
            minY = extent.yMinimum()
            maxY = extent.yMaximum()
            # 通过对齐最大值，确保裁剪范围完整
            maxX_aligned = maxX + pixel_width
            maxY_aligned = maxY + pixel_width
            minX_aligned = minX - pixel_width
            minY_aligned = minY - pixel_width
            # 设置 gdal_translate 命令和参数
            gdal_translate_cmd = [
                'gdal_translate',  # GDAL 工具名称
                '-projwin', str(minX_aligned), str(maxY_aligned), str(maxX_aligned), str(minY_aligned),
                # 裁剪范围（minX, maxY, maxX, minY）
                '-of', 'GTiff',  # 输出格式为 GeoTIFF
                input_raster,  # 输入栅格文件路径
                f"{output_raster}\\{key}.tif"  # 输出文件路径
            ]
            print(gdal_translate_cmd)
            # 运行命令
            try:
                subprocess.run(gdal_translate_cmd, check=True)
                print("gdal_translate 执行成功！")
                total = total + 1
            except subprocess.CalledProcessError as e:
                print(f"gdal_translate 执行失败: {e}")
        return {
            "状态": f"处理成功，共裁剪出{total}块栅格数据。"
        }