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

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

import os
import warnings

import netCDF4 as nc
import numpy as np
from osgeo import gdal, osr

'''
安装netCDF4库
conda install netCDF4 -c https://mirrors.aliyun.com/anaconda/pkgs/main/ -y
'''

warnings.filterwarnings('ignore')

def gis_NC_to_tif(in_data, out_path, ):
    """
    保存tif文件
    :param dtype: 值类型  float or int
    :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 + '\'')

    nc_data_obj = nc.Dataset(in_data)
    print(nc_data_obj.variables.keys())  # 变量名称列表
    # print(nc_data_obj.variables)  # 了解变量的基本信息

    # 这里要根据变量信息填写，不太的数据变量名称可能会不一致(高耦合)
    longitude = nc_data_obj.variables['x'][:]
    latitude = nc_data_obj.variables['y'][:]
    attr = np.asarray(nc_data_obj.variables['Band1'])  # 这里根据需求输入想要转换的波段名称

    # 影像的左上角和右下角坐标
    lon_min, lat_max, lon_max, lat_min = \
        [longitude.min(), latitude.max(), longitude.max(), latitude.min()]

    # 分辨率计算
    lat_array = len(latitude)
    lon_array = len(longitude)
    lon_res = (lon_max - lon_min) / (float(lon_array) - 1)
    lat_res = (lat_max - lat_min) / (float(lat_array) - 1)

    # 创建.tif文件
    driver = gdal.GetDriverByName('GTiff')  # 创建驱动
    dst_ds = driver.Create(out_path, lon_array, lat_array, 1, gdal.GDT_Float32)

    geotransform = (lon_min, lon_res, 0, lat_max, 0, -lat_res)  # 设置影像的显示范围
    dst_ds.SetGeoTransform(geotransform)
    # 定义投影（方法一:通过epsg来定义）
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)  # 4326-WGS_1984; 4490-GCS_China_Geodetic_Coordinate_System_2000
    dst_ds.SetProjection(srs.ExportToWkt())  # 给新建图层赋予投影信息
    # 定义投影（方法二:通过投影WKT编码来定义）
    # projection = 'GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0,AUTHORITY["EPSG","8901"]],UNIT["degree",0.0174532925199433,AUTHORITY["EPSG","9122"]],AXIS["Latitude",NORTH],AXIS["Longitude",EAST],AUTHORITY["EPSG","4326"]]'
    # dst_ds.SetProjection(projection)  # 给新建图层赋予投影信息

    # 去除异常值
    attr[attr[:, :] == -32768] = -99

    # 数据写出
    dst_ds.GetRasterBand(1).WriteArray(attr)
    dst_ds.GetRasterBand(1).SetNoDataValue(-99)
    dst_ds.FlushCache()  # 将数据写入硬盘
    del dst_ds  # 关闭tif文件


if __name__ == '__main__':

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