"""
@File    :   convertfile.py
@Time    :   2022/10/20 11:13:16
@Author  :   JuYongkang
@Version :   1.0
@Contact :   j_juyongkang@163.com
@Desc    :   nc文件转单一多波段tif文件
"""

# 模块导入
import numpy as np
import netCDF4 as nc
from osgeo import gdal, osr, ogr
import os


def NC_to_tiffs(file, Output_folder):
    """
    file: 待转化nc文件
    Output_folder: 转化tif文件目标文件夹
    """
    # 自动创建tif文件夹
    if not os.path.exists(Output_folder):
        os.mkdir(Output_folder)

    nc_data_obj = nc.Dataset(file)
    Lon = nc_data_obj.variables['lon'][:]
    Lat = nc_data_obj.variables['lat'][:]

    # 读取nc文件某个属性
    ndvi_arr = np.asarray(nc_data_obj.variables['data'])
    # 转化数据类型
    ndvi_arr_float = ndvi_arr.astype(float)

    # 查看数据维数
    print('**'*20)
    print(np.shape(ndvi_arr_float))
    print('**'*20)

    # 获取波段数
    # band = ndvi_arr_float.shape[0]
    band = ndvi_arr_float.shape[2]
    print('band为::', band)
    # 影像的左上角和右下角坐标
    LonMin, LatMax, LonMax, LatMin = [Lon.min(), Lat.max(), Lon.max(), Lat.min()]

    # 分辨率计算
    N_Lat = len(Lat)
    N_Lon = len(Lon)
    Lon_Res = (LonMax - LonMin) / (float(N_Lon) - 1)
    Lat_Res = (LatMax - LatMin) / (float(N_Lat) - 1)

    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif_name = Output_folder + '\\' + file.split('\\')[-1].split('.')[0] + '.tif'
    # out_tif = driver.Create(out_tif_name, N_Lon, N_Lat, band, gdal.GDT_Float32)
    out_tif = driver.Create(out_tif_name, N_Lat, N_Lon, band, gdal.GDT_Float32)

    # 设置影像的显示范围
    # -Lat_Res一定要是-的
    geotransform = (LonMin, Lon_Res, 0, LatMax, 0, -Lat_Res)
    out_tif.SetGeoTransform(geotransform)

    # 获取地理坐标系统信息，用于选取需要的地理坐标系统
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)  # 定义输出的坐标系为"WGS 84"，AUTHORITY["EPSG","4326"]
    out_tif.SetProjection(srs.ExportToWkt())  # 给新建图层赋予投影信息

    # 这里是使用np的swapaxes()方法转换轴的维度
    new_data = ndvi_arr_float.swapaxes(0, 2)
    new_data = new_data.swapaxes(1, 2)
    
    # 将数据写入到一个tif文件中
    print('band为::', band)
    for j in range(1, band+1, 1):
        data = new_data[j-1]
        print(np.shape(data))
        # 数据写出
        out_tif.GetRasterBand(j).WriteArray(data)  # 将数据写入内存，此时没有写入硬盘
    print('finish')
    out_tif.FlushCache()  # 将数据写入硬盘
    out_tif = None  # 注意必须关闭tif文件


if __name__ == '__main__':
    """
    file: 待转化nc文件
    target_path: 要保存tif文件的目录(是目录, 不是文件名)
    注: 转化后的tif文件和待转化nc文件文件名一致, 后缀不同
    """
    # 原nc文件地址
    file = r'C:\Users\JK\Desktop\wspd_12month_10.nc'
    # 转化tif文件目录
    target_path = r'C:\Users\JK\Documents\Projects\Test'

    NC_to_tiffs(file, target_path)