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

"""
@File    : 15_掩膜提取-批量.py
@Author  : fungis@163.com
@notice  :
"""
import os
from osgeo import gdal

filter_list = ['.TIF', '.tif']  # 设置过滤后的文件类型 当然可以设置多个类型


def find_filter_file(dirname):
    result = []  # 所有的文件
    for maindir, subdir, file_name_list in os.walk(dirname):
        # print("提示：当前主目录下的所有目录", "2:", subdir)  # 当前主目录下的所有目录
        # print("3:", file_name_list)  # 当前主目录下的所有文件

        for filename in file_name_list:
            apath = os.path.join(maindir, filename)  # 合并成一个完整路径
            ext = os.path.splitext(apath)[1]  # 获取文件后缀 [0]获取的是除了文件名以外的内容
            # 设置过滤器，过滤不符合格式的文件
            if ext in filter_list:
                result.append(apath)
    return result


if __name__ == '__main__':
    input_file_path = r'G:\pythonCode\python-package\VisualShowProject\study-geopandas\demo-ascMap\data-use\tif'  # 输入路径
    output_file_path = r'G:\pythonCode\python-package\VisualShowProject\study-geopandas\demo-ascMap\data-use\tifClip'  # 输出路径
    boundary_shp = r'G:\pythonCode\env\python-gdal-test\demo-time-axis\data-use\shp\Vic_polygon\Basin_boundary.shp'
    if not os.path.exists(output_file_path):
        # os.makedirs可以创建路径上多个
        os.makedirs(output_file_path)
    file_list = find_filter_file(input_file_path)
    if len(file_list) > 0:
        for raster_path in file_list:
            # 输出文件路径
            name = os.path.basename(raster_path).split('.')[0]
            result_file = output_file_path + os.sep + name + '.tif'  # 输出文件位置
            # clip
            gdal.Warp(result_file,  # 输出栅格
                      raster_path,  # 输入栅格
                      format='GTiff',
                      dstSRS='EPSG:4326',
                      cutlineDSName=boundary_shp,  # shp边界
                      cropToCutline=True,  # 按掩膜图层范围裁剪
                      outputType=gdal.GDT_Float32)
        print('------处理完成------')
    else:
        raise Exception('[Errno 1] 该路径下无{}类型文件'.format(filter_list))