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

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

import os

from qgis._core import QgsRasterFileWriter, QgsProcessingParameterFolderDestination
from qgis.core import (
    QgsVectorLayer,
    QgsRasterLayer,
    QgsFeature,
    QgsGeometry
)
import processing
from PyQt5.QtCore import QFileInfo
from qgis.analysis import QgsRasterCalculator, QgsRasterCalculatorEntry

from qgis.PyQt.QtCore import QCoreApplication
from qgis._core import QgsSpatialIndex, QgsField, QgsRectangle, QgsProcessingParameterVectorDestination, \
    QgsVectorFileWriter, QgsWkbTypes, QgsProcessingParameterNumber, QgsProcessingParameterBoolean, QgsGeometry, \
    QgsFeature
from qgis.core import (QgsProcessing,
                       QgsVectorLayer,
                       QgsFeatureSink,
                       QgsProcessingException,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFile,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterFeatureSink)
from qgis import processing


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'
        ))
        # 通过 setExtensions 方法设置多个扩展名
        # self.parameter(self.INPUT_RASTER).setExtensions(['tif', 'img'])  # 支持 .shp, .csv, .txt
        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路径
        # input_shp = "E:\\projects\\遥感技术部需求\\裁剪栅格\\data\\grid_extent.shp"
        # 输入栅格数据路径
        # input_raster = "D:\\gisdata\\tif\\unnamed_卫图_Level_12.tif"
        # 输出裁剪后的栅格路径
        # output_raster = "E:\\projects\\遥感技术部需求\\裁剪栅格\\output\\"
        # 加载输入Shapefile和栅格数据
        vector_layer = QgsVectorLayer(input_shp, "vector_layer", "ogr")
        raster_layer = QgsRasterLayer(input_raster, "raster_layer")
        fields = vector_layer.fields()
        # 遍历输入图层的所有要素，并为每个要素保存一个单独的Shapefile文件
        total = 0
        for idx, feature in enumerate(vector_layer.getFeatures()):
            # 获取当前要素的几何数据
            geometry = feature.geometry()
            key = feature["新图号"]

            # 创建一个新的 Shapefile 文件名，可以使用索引或 ID 来命名
            output_shp = os.path.join(output_raster, f"{key}.shp")

            # 创建输出shapefile
            writer = QgsVectorFileWriter(output_shp, 'UTF-8', fields, QgsWkbTypes.MultiPolygon, vector_layer.crs(),
                                         'ESRI Shapefile')
            writer.addFeature(feature)

            # 完成写入并关闭输出文件
            del writer

            # 使用QGIS的processing模块进行栅格裁剪
            params = {
                'INPUT': raster_layer,  # 输入栅格
                'MASK': output_shp,  # 用于裁剪的矢量图层
                'OUTPUT': f"{output_raster}\\{key}.tif",  # 输出裁剪后的栅格文件
                'CROP_TO_CUTLINE': True
            }

            # 执行裁剪操作
            processing.run('gdal:cliprasterbymasklayer', params)
            print(f"Cropped raster saved to {output_raster}")
            total = total + 1
        return {
            "状态": f"处理成功，共裁剪出{total}块栅格数据。"
        }