"""
@Project ：microproduct 
@File ：DEMJoint
@Function ：主函数
@Author ：LMM
@Date ：2021/10/19 14:39
@Version ：1.0.0
"""
from osgeo import gdal, osr
import os
import numpy as np


class DEMProcess:
    """
    DEM拼接、重采样
    """
    def __init__(self):
        pass

    @staticmethod
    def get_extent(fn):
        '''
        原文链接：https://blog.csdn.net/XBR_2014/article/details/85255412
        '''
        ds = gdal.Open(fn)
        rows = ds.RasterYSize
        cols = ds.RasterXSize
        # 获取图像角点坐标
        gt = ds.GetGeoTransform()
        minx = gt[0]
        maxy = gt[3]
        maxx = gt[0] + gt[1] * rows
        miny = gt[3] + gt[5] * cols
        return (minx, maxy, maxx, miny)

    @staticmethod
    def img_mosaic(in_files, out_dem_path):
        # 通过两两比较大小,将最终符合条件的四个角点坐标保存
        # 即为拼接图像的四个角点坐标
        minX, maxY, maxX, minY = DEMProcess.get_extent(in_files[0])
        for fn in in_files[1:]:
            minx, maxy, maxx, miny = DEMProcess.get_extent(fn)
            minX = min(minX, minx)
            maxY = max(maxY, maxy)
            maxX = max(maxX, maxx)
            minY = min(minY, miny)

        # 获取输出图像的行列数
        in_ds = gdal.Open(in_files[0])
        bands_num = in_ds.RasterCount
        gt = in_ds.GetGeoTransform()
        rows = int((maxX - minX) / abs(gt[5]))
        cols = int((maxY - minY) / gt[1])

        # 判断栅格数据的数据类型
        datatype = gdal.GDT_UInt16

        # 创建输出图像
        driver = gdal.GetDriverByName('GTiff')
        out_dem = os.path.join(out_dem_path, 'mosaic0.tif')
        out_ds = driver.Create(out_dem, cols, rows, bands_num, datatype)
        out_ds.SetProjection(in_ds.GetProjection())

        gt = list(in_ds.GetGeoTransform())
        gt[0], gt[3] = minX, maxY
        out_ds.SetGeoTransform(gt)

        for fn in in_files:
            in_ds = gdal.Open(fn)
            x_size = in_ds.RasterXSize
            y_size = in_ds.RasterYSize
            trans = gdal.Transformer(in_ds, out_ds, [])
            success, xyz = trans.TransformPoint(False, 0, 0)
            x, y, z = map(int, xyz)
            for i in range(1, bands_num + 1):
                data = in_ds.GetRasterBand(i).ReadAsArray()
                out_band = out_ds.GetRasterBand(i)
                out_data = out_band.ReadAsArray(x, y, x_size, y_size)
                data = np.maximum(data, out_data)
                out_band.WriteArray(data, x, y)

        del in_ds, out_band, out_ds

    @staticmethod
    def dem_clip(OutFilePath, DEMFilePath, SelectArea):
        '''
        根据选择范围裁剪DEM,并输出
        agrs:
            outFilePath:裁剪DEM输出地址
            DEMFilePath:被裁减DEM地址
            SelectArea:list  [(xmin,ymax),(xmax,ymin)] 框选范围 左上角，右下角
        '''
        DEM_ptr = gdal.Open(DEMFilePath)
        DEM_GeoTransform = DEM_ptr.GetGeoTransform()  # 读取影像的投影变换
        DEM_InvGeoTransform = gdal.InvGeoTransform(DEM_GeoTransform)
        SelectAreaArrayPoints = [gdal.ApplyGeoTransform(DEM_InvGeoTransform, p[0], p[1]) for p in SelectArea]
        SelectAreaArrayPoints = list(map(lambda p: (int(p[0]), int(p[1])), SelectAreaArrayPoints))  # 确定坐标

        [(ulx, uly), (brx, bry)] = SelectAreaArrayPoints
        rowCount, colCount = bry - uly, brx - ulx

        # 输出DEM的桌面坐标转换
        Out_Transfrom = list(DEM_GeoTransform)
        Out_Transfrom[0] = SelectArea[0][0]
        Out_Transfrom[3] = SelectArea[0][1]

        # 构建输出DEM
        Bands_num = DEM_ptr.RasterCount
        gtiff_driver = gdal.GetDriverByName('GTiff')
        datatype = gdal.GDT_UInt16
        out_dem = gtiff_driver.Create(OutFilePath, colCount, rowCount, Bands_num, datatype)
        out_dem.SetProjection(DEM_ptr.GetProjection())
        out_dem.SetGeoTransform(Out_Transfrom)

        for i in range(1, Bands_num + 1):
            data_band = DEM_ptr.GetRasterBand(i)
            out_band = out_dem.GetRasterBand(i)
            data = data_band.ReadAsArray(ulx, uly, colCount, rowCount)
            out_band.WriteArray(data)
        del out_dem

    @staticmethod
    def dem_resample(in_dem_path, out_dem_path):
        '''
        DEM重采样函数，默认坐标系为WGS84
        agrs:
            in_dem_path: 输入的DEM文件夹路径
            meta_file_path: 输入的xml元文件路径
            out_dem_path: 输出的DEM文件夹路径
        '''
        # 读取文件夹中所有的DEM
        dem_file_paths=[os.path.join(in_dem_path,dem_name) for dem_name in os.listdir(in_dem_path) if dem_name.find(".tif")>=0 and dem_name.find(".tif.")==-1]
        spatialreference=osr.SpatialReference()
        spatialreference.SetWellKnownGeogCS("WGS84")  # 设置地理坐标，单位为度 degree # 设置投影坐标，单位为度 degree
        spatialproj=spatialreference.ExportToWkt()  # 导出投影结果
        # 将DEM拼接成一张大图
        mergeFile =gdal.BuildVRT(os.path.join(out_dem_path,"mergeDEM.tif"), dem_file_paths)
        out_DEM=os.path.join(out_dem_path,"mosaic.tif")
        gdal.Warp(out_DEM,
                    mergeFile,
                    format="GTiff",
                    dstSRS=spatialproj,
                    dstNodata=-9999,
                    outputType=gdal.GDT_Float32)
        return out_DEM


# if __name__ == "__main__":
#     DEMProcess = DEMProcess()
#     in_dem_path = r'F:\大气延迟\out_dem'
#     out_dem_path = r'F:\大气延迟\out_dem'
#     DEMProcess.dem_resample(in_dem_path, out_dem_path)