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

"""
@File    : 08_文件转换之HDF转TIF.py
@Author  : fungis@163.com
@notice  :
"""

import os

from osgeo import gdal, gdalconst


def gis_hdf_to_tif(in_data, out_path, dtype):
    """
    保存tif文件
    :param in_data: 原始数据路径
    :param out_path: 保存tif文件路径
    :param dtype: gdal数据类型, 默认自动根据输入识别
    :return:
    """
    if os.path.exists(in_data) is False:
        raise Exception('[Errno 2] 该文件不存在: \'' + in_data + '\'')

    dataset = gdal.Open(in_data, gdalconst.GA_ReadOnly)  # 读取遥感数据信息
    band_count = dataset.RasterCount  # 波段数
    if band_count > 0:
        hdf_to_tif_write(dataset, band_count, dtype, out_path)
    else:
        raise Exception('[Errno 4] 该文件波段异常: \'' + in_data + '\'')


def hdf_to_tif_write(dataset, band_count, dtype, out_path):
    geo_transform = dataset.GetGeoTransform()  # 获取地理空间信息
    columns = dataset.RasterXSize
    row = dataset.RasterYSize
    dim_z = band_count  # 波段数

    # 创建.tif文件
    driver = gdal.GetDriverByName('GTiff')  # 创建驱动
    # 创建文件
    if dtype == "float" or dtype == "float32":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_Float32)
    elif dtype == "float64":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_Float64)
    elif dtype == "int":
        dst_ds = driver.Create(out_path, columns, row, dim_z, gdal.GDT_UInt16)
    else:
        raise Exception('[Errno 4] 未设置值类型或设置字符串存在错误！')
    dst_ds.SetGeoTransform(geo_transform)  # 设置几何信息
    dst_ds.SetProjection(dataset.GetProjection())  # 设置投影
    for i in range(1, (int(band_count) + 1)):
        band = dataset.GetRasterBand(i)
        cell_map = band.ReadAsArray()
        dst_ds.GetRasterBand(i).WriteArray(cell_map)  # 写入像元信息
        dst_ds.FlushCache()  # 写入硬盘
    del dst_ds  # 关闭tif文件


if __name__ == '__main__':
    # raster_path = r'G:\temp\hdf\MOSAIC_TMP_2019273.hdf'  # 输入你的hdf文件(具体路径)
    raster_path = r'./data-use/hdf/MOSAIC_TMP_2019273.hdf'  # 输入你的hdf文件(相对路径)
    # 输出文件路径
    result_file = r'./results/hdf/test20221012.10.tif'
    # 创建文件存放文件夹
    dir_path = result_file[:result_file.rindex('/')]
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    # 执行转换函数
    gis_hdf_to_tif(raster_path, result_file, 'float')
    print('------处理完成------')