#!/Users/christmas/opt/anaconda3/bin/python3
# -*- coding: utf-8 -*-
#  日期 : 2023/9/18 15:38
#  作者 : Christmas
#  邮箱 : 273519355@qq.com
#  项目 : Project
#  版本 : python 3
#  摘要 :
"""

"""
from typing import Union

import numpy as np
import xarray as xr

from .. import new_filename

__all__ = ['add_coors_attrs', 'add_lon_attrs', 'add_lat_attrs', 'add_time_attrs', 'add_TIME_attrs',
           'add_wind_attrs', 'add_temperature_attrs', 'add_precipitation_attrs', 'add_slp_attrs',
           'add_wave_attrs', 'add_ice_attrs', 'add_visibility_attrs',
           'fixed_encoding', 'calc_scale_offset', 'compress_scale_offset',
           ]


def add_coors_attrs(_ds):
    """
    给xarray.Dataset的坐标添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    lon_name = next((x for x in ['lon', 'longitude'] if x in _ds.dims), None)
    lat_name = next((x for x in ['lat', 'latitude'] if x in _ds.dims), None)
    time_name = next((x for x in ['time'] if x in _ds.dims), None)
    TIME_name = next((x for x in ['TIME'] if x in _ds.dims), None)

    add_lon_attrs(_ds, lon_name)
    add_lat_attrs(_ds, lat_name)
    add_time_attrs(_ds, time_name)
    add_TIME_attrs(_ds, TIME_name, dim_name='DateStr')
    return _ds
    

def add_lon_attrs(_ds, lon_name):
    """
    给经度添加属性
    :param _ds: xarray.Dataset
    :param lon_name: 经度维度名
    :return: xarray.Dataset
    """
    _ds[lon_name].attrs = {
        'units': 'degrees_east',
        'long_name': 'longitude',
        'standard_name': 'longitude',
        'axis': 'X',
        'westernmost': np.round(_ds[lon_name].data.min(), 2),
        'easternmost': np.round(_ds[lon_name].data.max(), 2),
    }
    _ds[lon_name].encoding = {'_FillValue': None}
    return _ds


def add_lat_attrs(_ds, lat_name):
    """
    给纬度添加属性
    :param _ds: xarray.Dataset
    :param lat_name: 纬度维度名
    :return: xarray.Dataset
    """
    _ds[lat_name].attrs = {
        'units': 'degrees_north',
        'long_name': 'latitude',
        'standard_name': 'latitude',
        'axis': 'Y',
        'southernmost': np.round(_ds[lat_name].data.min(), 2),
        'northernmost': np.round(_ds[lat_name].data.max(), 2)
    }
    _ds[lat_name].encoding = {'_FillValue': None}
    return _ds


def add_time_attrs(_ds, time_name):
    """
    给时间添加属性
    :param _ds: xarray.Dataset
    :param time_name: 时间维度名
    :return: xarray.Dataset
    """
    _ds[time_name].attrs = {
        'units': 'seconds since 1970-01-01 00:00:00',
        'long_name': 'UTC time',
        'standard_name': 'UTC_time',
        'axis': 'T',
        'calendar': 'gregorian',
    }
    _ds[time_name].encoding = {'_FillValue': None, 'dtype': 'float64', 'unlimited_dims': 'time'}
    return _ds


def add_TIME_attrs(_ds, TIME_name, dim_name='DateStr'):
    """
    给TIME添加属性
    :param _ds: xarray.Dataset
    :param TIME_name: TIME名
    :param dim_name: TIME维度名
    :return: xarray.Dataset
    """
    _ds[TIME_name].attrs = {
        'reference_time': ''.join(_ds[TIME_name].data[0].astype(str))[:10],
        'long_name': 'UTC time',
        'standard_name': 'UTC_time',
        'calendar': 'gregorian',
        'start_time': ''.join(_ds[TIME_name].data[0].astype(str)),
        'end_time': ''.join(_ds[TIME_name].data[-1].astype(str))
    }
    _ds[TIME_name].encoding['char_dim_name'] = dim_name
    return _ds


def add_wind_attrs(_ds):
    """
    给xarray.Dataset的风添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    U10_name = next((x for x in ['wind_U10', 'U10', 'u10'] if x in _ds.data_vars), None)
    V10_name = next((x for x in ['wind_V10', 'V10', 'v10'] if x in _ds.data_vars), None)
    _ds[U10_name].attrs = {
        'units': 'm s-1',
        'long_name': '10 meter U wind component',
        'standard_name': 'eastward_wind',
    }
    _ds[V10_name].attrs = {
        'units': 'm s-1',
        'long_name': '10 meter V wind component',
        'standard_name': 'northward_wind',
    }
    _ds[U10_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
    _ds[V10_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if 'AGL_heights' in _ds.variables:
        _ds['AGL_heights'].attrs = {
            'units': 'm',
            'long_name': 'Above Ground Level Height',
            'standard_name': 'height',
        }
        _ds['AGL_heights'].encoding = {'_FillValue': None}

    if 'U_AGL' in _ds.variables and 'V_AGL' in _ds.variables:
        _ds['U_AGL'].attrs = {
            'units': 'm s-1',
            'long_name': 'U-wind at AGL',
            'standard_name': 'eastward_wind',
        }
        _ds['V_AGL'].attrs = {
            'units': 'm s-1',
            'long_name': 'V-wind at AGL',
            'standard_name': 'northward_wind',
        }
        _ds['U_AGL'].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
        _ds['V_AGL'].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if 'MSL_heights' in _ds.variables:
        _ds['MSL_heights'].attrs = {
            'units': 'm',
            'long_name': 'Mean Sea Level Height',
            'standard_name': 'height',
        }
        _ds['MSL_heights'].encoding = {'_FillValue': None}

    if 'U_MSL' in _ds.variables and 'V_MSL' in _ds.variables:
        _ds['U_MSL'].attrs = {
            'units': 'm s-1',
            'long_name': 'U-wind at MSL',
            'standard_name': 'eastward_wind',
        }
        _ds['V_MSL'].attrs = {
            'units': 'm s-1',
            'long_name': 'V-wind at MSL',
            'standard_name': 'northward_wind',
        }
        _ds['U_MSL'].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
        _ds['V_MSL'].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if 'WSPD10MAX' in _ds.variables:
        _ds['WSPD10MAX'].attrs = {
            'units': 'm s-1',
            'long_name': 'Maximum 10m wind speed',
            'standard_name': 'maximum_10m_wind_speed',
        }
        _ds['WSPD10MAX'].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    return _ds


def add_temperature_attrs(_ds):
    """
    给xarray.Dataset的温度添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    T2_name = next((x for x in ['temperature', 'T2', 't2'] if x in _ds.data_vars), None)

    _ds[T2_name].attrs = {
        'units': 'K',
        'long_name': 'temperature at 2m',
        'standard_name': 'temperature_at_2m'
    }
    _ds[T2_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
    return _ds


def add_visibility_attrs(_ds):
    """
    给xarray.Dataset的能见度添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    visibility_name = next((x for x in ['visibility', 'AFWA_VIS', 'VISIBILITY'] if x in _ds.data_vars), None)

    _ds[visibility_name].attrs = {
        'units': 'm',
        'long_name': 'visibility',
        'standard_name': 'visibility'
    }
    _ds[visibility_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
    return _ds


def add_precipitation_attrs(_ds):
    """
    给xarray.Dataset的降水添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    precipitation_name = next((x for x in ['precipitation', 'tp', 'TP', 'Precipitation'] if x in _ds.data_vars), None)

    _ds[precipitation_name].attrs = {
        'units': 'mm',
        'long_name': 'precipitation, positive for ocean gaining water',
        'standard_name': 'precipitation_positive_ocean_gaining_water'
    }
    _ds[precipitation_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
    return _ds


def add_slp_attrs(_ds):
    """
    给xarray.Dataset的海平面气压添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    slp_name = next((x for x in ['slp', 'SLP'] if x in _ds.data_vars), None)

    _ds[slp_name].attrs = {
        'units': 'Pa',
        'long_name': 'sea level pressure',
        'standard_name': 'sea_level_pressure'
    }
    _ds[slp_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)
    return _ds


def add_wave_attrs(_ds):
    """
    给xarray.Dataset的海浪添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    swh_name = next((x for x in ['swh', 'SWH'] if x in _ds.data_vars), None)  # 有效波高
    mwp_name = next((x for x in ['mwp', 'MWP'] if x in _ds.data_vars), None)  # 平均波周期
    mwd_name = next((x for x in ['mwd', 'MWD'] if x in _ds.data_vars), None)  # 平均波向
    mwl_name = next((x for x in ['mwl', 'MWL'] if x in _ds.data_vars), None)  # 平均波长

    shww_name = next((x for x in ['shww', 'SHWW'] if x in _ds.data_vars), None)  # 风浪有效波高
    shts_name = next((x for x in ['shts', 'SHTS'] if x in _ds.data_vars), None)  # 涌浪有效波高
    mdww_name = next((x for x in ['mdww', 'MDWW'] if x in _ds.data_vars), None)  # 风浪平均波向
    mdts_name = next((x for x in ['mdts', 'MDTS'] if x in _ds.data_vars), None)  # 涌浪平均波向
    mpww_name = next((x for x in ['mpww', 'MPWW'] if x in _ds.data_vars), None)  # 风浪平均波周期
    mpts_name = next((x for x in ['mpts', 'MPTS'] if x in _ds.data_vars), None)  # 涌浪平均波周期
    
    pp1d_name = next((x for x in ['pp1d', 'PP1D'] if x in _ds.data_vars), None)  # 峰值周期
    hmax_name = next((x for x in ['hmax', 'HMAX'] if x in _ds.data_vars), None)  # 最大波高

    if swh_name in _ds:
        _ds[swh_name].attrs = {
            'units': 'm',
            'long_name': 'significant wave height',
            'standard_name': 'significant_wave_height'
        }
        _ds[swh_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mwp_name in _ds:
        _ds[mwp_name].attrs = {
            'units': 's',
            'long_name': 'mean wave period',
            'standard_name': 'mean_wave_period'
        }
        _ds[mwp_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mwd_name in _ds:
        _ds[mwd_name].attrs = {
            'units': 'degrees',
            'long_name': 'mean wave direction',
            'standard_name': 'mean_wave_direction',
            'direction_0': 'coming from the north',
            'direction_90': 'coming from the east'
        }
        _ds[mwd_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mwl_name in _ds:
        _ds[mwl_name].attrs = {
            'units': 'm',
            'long_name': 'mean wave length',
            'standard_name': 'mean_wave_length'
        }
        _ds[mwl_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if pp1d_name in _ds:
        _ds[pp1d_name].attrs = {
            'units': 's',
            'long_name': 'peak period',
            'standard_name': 'peak_period'
        }
        _ds[pp1d_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if shww_name in _ds:
        _ds[shww_name].attrs = {
            'units': 'm',
            'long_name': 'significant height of wind waves',
            'standard_name': 'significant_height_of_wind_waves'
        }
        _ds[shww_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if shts_name in _ds:
        _ds[shts_name].attrs = {
            'units': 'm',
            'long_name': 'significant height of total swell',
            'standard_name': 'significant_height_of_total_swell'
        }
        _ds[shts_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mdww_name in _ds:
        _ds[mdww_name].attrs = {
            'units': 'degrees',
            'long_name': 'mean direction of wind waves',
            'standard_name': 'mean_direction_of_wind_waves',
            'direction_0': 'coming from the north',
            'direction_90': 'coming from the east'
        }
        _ds[mdww_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mdts_name in _ds:
        _ds[mdts_name].attrs = {
            'units': 'degrees',
            'long_name': 'mean direction of total swell',
            'standard_name': 'mean_direction_of_total_swell',
            'direction_0': 'coming from the north',
            'direction_90': 'coming from the east'
        }
        _ds[mdts_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mpww_name in _ds:
        _ds[mpww_name].attrs = {
            'units': 's',
            'long_name': 'mean period of wind waves',
            'standard_name': 'mean_period_of_wind_waves'
        }
        _ds[mpww_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if mpts_name in _ds:
        _ds[mpts_name].attrs = {
            'units': 's',
            'long_name': 'mean period of total swell',
            'standard_name': 'mean_period_of_total_swell'
        }
        _ds[mpts_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if hmax_name in _ds:
        _ds[hmax_name].attrs = {
            'units': 'm',
            'long_name': 'maximum wave height',
            'standard_name': 'maximum_wave_height'
        }
        _ds[hmax_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    return _ds


def add_ice_attrs(_ds):
    """
    给xarray.Dataset的海冰添加属性
    :param _ds: xarray.Dataset
    :return: xarray.Dataset
    """
    aice_name = next((x for x in ['aice'] if x in _ds.data_vars), None)
    tice_name = next((x for x in ['tice'] if x in _ds.data_vars), None)

    if aice_name in _ds:
        _ds[aice_name].attrs = {
            'units': 'fraction',
            'long_name': 'sea ice concentration',
            'standard_name': 'sea_ice_concentration'
        }
        _ds[aice_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    if tice_name in _ds:
        _ds[tice_name].attrs = {
            'units': 'm',
            'long_name': 'sea ice thickness',
            'standard_name': 'sea_ice_thickness'
        }
        _ds[tice_name].encoding = dict(_FillValue=9.962100e+36, dtype='single', zlib=True)

    return _ds


def fixed_encoding(_ds, **kwargs):
    """
    修正xarray.Dataset的encoding
    """
    _ds.encoding['unlimited_dims'] = kwargs.get('unlimited_dims', 'time')
    _ds.encoding['format'] = kwargs.get('format', 'NETCDF4')
    _ds.encoding['engine'] = kwargs.get('engine', 'h5netcdf')
    _ds.encoding['dtype'] = kwargs.get('dtype', 'float32')
    _ds.encoding['complevel'] = kwargs.get('complevel', 4)
    _ds.encoding['zlib'] = kwargs.get('zlib', True)
    _ds.encoding['shuffle'] = kwargs.get('shuffle', True)
    _ds.encoding['fletcher32'] = kwargs.get('fletcher32', True)
    _ds.encoding['contiguous'] = kwargs.get('contiguous', True)
    _ds.encoding.update(kwargs)
    return _ds


def calc_scale_offset(
        _data: np.ndarray,
        _dtype=np.int16
) -> tuple[float, float]:
    """
    计算scale_factor和add_offset
    :param _data: 数据数组
    :param _dtype: 压缩后的数据类型，默认为int16
    :return: scale_factor, add_offset
    """
    valid_min = np.nanmin(_data)
    valid_max = np.nanmax(_data)

    if valid_min == valid_max:
        scale_factor = 1.0
        add_offset = valid_min
    else:
        scale_factor = (valid_max - valid_min) / (np.iinfo(_dtype).max - np.iinfo(_dtype).min)
        add_offset = valid_min - np.iinfo(_dtype).min * scale_factor

    return scale_factor, add_offset


def compress_scale_offset(
        _ds: xr.Dataset,
        _var_name: Union[str, list[str]] = None,
        _dtype=np.int16
) -> xr.Dataset:
    """
    给xarray.Dataset的变量添加压缩属性
    :param _ds: xarray.Dataset
    :param _var_name: 变量名或变量名列表
    :param _dtype: 压缩后的数据类型，默认为int16
    :return: xarray.Dataset
    """

    if _var_name is None:
        _var_name = list(_ds.data_vars.keys())
    elif isinstance(_var_name, str):
        _var_name = [_var_name]

    for var in _var_name:
        if var in _ds.data_vars:
            data = _ds[var].data
            scale_factor, add_offset = calc_scale_offset(data, _dtype=_dtype)

            _ds[var].encoding['dtype'] = _dtype
            _ds[var].encoding['scale_factor'] = scale_factor
            _ds[var].encoding['add_offset'] = add_offset
            _ds[var].encoding['_FillValue'] = np.iinfo(_dtype).min

    return _ds
