
import numpy as np
import math
import meteva
import os
import traceback
import warnings
warnings.filterwarnings("ignore")

def write_griddata_to_micaps4(da,save_path = "a.txt",creat_dir = False,effectiveNum = 3,show = False,title = None):
    """
    输出micaps4格式文件
    :param da:xarray多维数据信息
    :param path:存储路径
    :param effectiveNum 有效数字 默认：3
    :return 最终按照需要保存的路径，将da数据保存为m4格式
    """
    try:
        dir = os.path.split(os.path.abspath(save_path))[0]
        if not os.path.isdir(dir):
            if not creat_dir:
                print("文件夹："+dir+"不存在")
                return False
            else:
                meteva.base.tool.path_tools.creat_path(save_path)

        grid = meteva.base.basicdata.get_grid_of_data(da)
        nlon = grid.nlon
        nlat = grid.nlat
        slon = grid.slon
        slat = grid.slat
        elon = grid.elon
        elat = grid.elat
        dlon = grid.dlon
        dlat = grid.dlat
        level = grid.levels[0]
        stime = grid.stime_str
        year = stime[0:4]
        month = stime[4:6]
        day = stime[6:8]
        hour = stime[8:10]
        hour_range = str(grid.dtimes[0])
        values = da.values
        grid_values = np.squeeze(values)
        vmax = math.ceil(max(grid_values.flatten()))
        vmin = math.ceil(min(grid_values.flatten()))

        dif = (vmax - vmin) / 10.0
        if dif ==0:
            inte = 1
        else:
            inte = math.pow(10, math.floor(math.log10(dif)))
        # 用基本间隔，将最大最小值除于间隔后小数点部分去除，最后把间隔也整数化
        r = dif / inte
        if r < 3 and r >= 1.5:
            inte = inte * 2
        elif r < 4.5 and r >= 3:
            inte = inte * 4
        elif r < 5.5 and r >= 4.5:
            inte = inte * 5
        elif r < 7 and r >= 5.5:
            inte = inte * 6
        elif r >= 7:
            inte = inte * 8
        vmin = inte * ((int)(vmin / inte) - 1)
        vmax = inte * ((int)(vmax / inte) + 1)

        end = len(save_path)
        start = max(0, end - 16)

        if title is None:
            title = ("diamond 4 " + save_path[start:end] + "\n"
                     + year + " " + month + " " + day + " " + hour + " " + hour_range + " " + str(level) + "\n"
                     + "{:.6f}".format(grid.dlon) + " " + "{:.6f}".format(grid.dlat) + " " + "{:.6f}".format(grid.slon) + " " +"{:.6f}".format(grid.elon) + " "
                     + "{:.6f}".format(grid.slat) + " " + "{:.6f}".format(grid.elat) + " " + str(grid.nlon) + " " + str(grid.nlat) + " "
                     + str(inte) + " " + str(vmin) + " " + str(vmax) + " 1 0")
        else:

            title = ("diamond 4 "+ title +"\n"
                     + year + " " + month + " " + day + " " + hour + " " + hour_range + " " + str(level) + "\n"
                     + "{:.6f}".format(grid.dlon) + " " + "{:.6f}".format(grid.dlat) + " " + "{:.6f}".format(grid.slon) + " " + "{:.6f}".format(grid.elon) + " "
                     + "{:.6f}".format(grid.slat) + " " + "{:.6f}".format(grid.elat) + " " + str(grid.nlon) + " " + str(grid.nlat) + " "
                     + str(inte) + " " + str(vmin) + " " + str(vmax) + " 1 0")

        # 二维数组写入micaps文件
        format_str = "%." + str(effectiveNum) + "f "

        np.savetxt(save_path, grid_values, delimiter=' ',
                   fmt=format_str, header=title, comments='')
        if show:
            print('成功输出至'+ save_path)
        return True
    except:
        exstr = traceback.format_exc()
        print(exstr)
        return False

def write_griddata_to_nc(da,save_path = "a.txt",creat_dir = False,effectiveNum = 3,show = False):
    try:
        dir = os.path.split(os.path.abspath(save_path))[0]
        if not os.path.isdir(dir):
            if not creat_dir:
                print("文件夹："+dir+"不存在")
                return False
            else:
                meteva.base.tool.path_tools.creat_path(save_path)
        scale_factor = math.pow(10,-effectiveNum)
        #print(scale_factor)
        encodingdict = {da.name:{
                            'dtype': 'int32',
                            'scale_factor': scale_factor,
                             'zlib': True,
                            '_FillValue':None
                            }
                        }

        da.to_netcdf(save_path,encoding = encodingdict)
        if show:
            print('成功输出至' + save_path)
        return True
    except:
        exstr = traceback.format_exc()
        print(exstr)
        return False
def write_griddata_to_micaps11(wind,save_path = "a.txt",creat_dir = False,effectiveNum = 3,show = False,title = None):
    try:
        dir = os.path.split(os.path.abspath(save_path))[0]
        if not os.path.isdir(dir):
            if not creat_dir:
                print("文件夹："+dir+"不存在")
                return False
            else:
                meteva.base.tool.path_tools.creat_path(save_path)
        grid0 = meteva.base.basicdata.get_grid_of_data(wind)
        nlon = grid0.nlon
        nlat = grid0.nlat
        slon = grid0.slon
        slat = grid0.slat
        elon = grid0.elon
        elat = grid0.elat
        dlon = grid0.dlon
        dlat = grid0.dlat
        level = grid0.levels[0]
        stime = grid0.stime_str
        year = stime[0:4]
        month = stime[4:6]
        day = stime[6:8]
        hour = stime[8:10]
        values = wind.values
        grid_values = np.squeeze(values).reshape(2*nlat,nlon)

        end = len(save_path)
        start = max(0, end - 16)

        if title is None:
            title = ("diamond 11 " + save_path[start:end] + "\n"
                     + year + " " + month + " " + day + " " + hour + " " + str(level) + "\n"
                     + str(grid0.dlon) + " " + str(grid0.dlat) + " " + str(grid0.slon) + " " + str(grid0.elon) + " "
                     + str(grid0.slat) + " " + str(grid0.elat) + " " + str(grid0.nlon) + " " + str(grid0.nlat))
        else:

            title  = ("diamond 11 " + title+ "\n"
                    +year + " "+ month + " " + day+ " " +hour +" " + str(level)+"\n"
                    + str(grid0.dlon) + " " + str(grid0.dlat) + " " + str(grid0.slon) + " " + str(grid0.elon) + " "
                    + str(grid0.slat) + " " + str(grid0.elat) + " " + str(grid0.nlon) + " " + str(grid0.nlat))

        format_str = "%." + str(effectiveNum) + "f "

        np.savetxt(save_path, grid_values, delimiter=' ',
                   fmt=format_str, header=title, comments='')
        if show:
            print('成功输出至' + save_path)
        return True
    except:
        exstr = traceback.format_exc()
        print(exstr)
        return False

def tran_griddata_to_gds_flow(da):
    grid0 = meteva.base.get_grid_of_data(da)
    discriminator = b"mdfs"
    if len(grid0.members)==1:
        data_type = 4
    elif len(grid0.members)==2:
        data_type =11
    else:
        print("仅支持micap4类和micaps11类数据输出成GDS格式")
    data_type_byte = np.ndarray.tobytes(np.array([data_type]).astype(np.int16))
    mName = grid0.members[0]
    mName = mName.encode(encoding='utf-8')
    if len(mName)<20:
        mName = mName + np.ndarray.tobytes(np.zeros(20-len(mName)).astype(np.int8))

    eleName = b""
    if "eleName" in da.attrs.keys():
        eleName = da.attrs["eleName"]
    if len(eleName) < 50:
        eleName = eleName + np.ndarray.tobytes(np.zeros(50 - len(eleName)).astype(np.int8))

    description =b""
    if "description" in da.attrs.keys():
        description = da.attrs["description"]
    if len(description) < 30:
        description = description + np.ndarray.tobytes(np.zeros(30 - len(description)).astype(np.int8))


    level = np.ndarray.tobytes(np.array(grid0.levels[0]).astype(np.float32))
    y_m_d_h_timezone_peroid = np.ndarray.tobytes(np.array([2021,1,1,8,8,0]).astype(np.int32))
    slon_elon_dlon = np.ndarray.tobytes(np.array([grid0.slon,grid0.elon,grid0.dlon]).astype(np.float32))
    nlon = np.ndarray.tobytes(np.array([grid0.nlon]).astype(np.int32))
    slat_elat_dlat = np.ndarray.tobytes(np.array([grid0.slat, grid0.elat, grid0.dlat]).astype(np.float32))
    nlat  =  np.ndarray.tobytes(np.array([grid0.nlat]).astype(np.int32))
    vmin,vmax,inte = meteva.base.tool.plot_tools.get_isoline_set(da)
    sValue_eValue_dValue = np.ndarray.tobytes(np.array([vmin,vmax,inte]).astype(np.float32))
    blank = np.ndarray.tobytes(np.zeros(100).astype(np.int8))
    value_bytes = np.ndarray.tobytes(da.values.astype(np.float32))
    bytes1 = discriminator+data_type_byte+mName+eleName+description+level+y_m_d_h_timezone_peroid
    bytes2 = slon_elon_dlon+nlon+slat_elat_dlat+nlat+sValue_eValue_dValue+blank+value_bytes
    bytes = bytes1+bytes2

    return bytes

def write_griddata_to_gds_file(da,save_path = "a.txt",creat_dir = False,show = False):
    try:
        dir = os.path.split(os.path.abspath(save_path))[0]
        if not os.path.isdir(dir):
            if not creat_dir:
                print("文件夹："+dir+"不存在")
                return False
            else:
                meteva.base.tool.path_tools.creat_path(save_path)

        bytes = tran_griddata_to_gds_flow(da)
        br = open(save_path, 'wb')
        br.write(bytes)
        br.close()
        if show:
            print('成功输出至' + save_path)
        return True
    except:
        exstr = traceback.format_exc()
        print(exstr)
        return False

if __name__ == "__main__":
    grd = meteva.base.read_griddata_from_micaps4(r"H:\test_data\input\meb\m4.txt")
    write_griddata_to_gds_file(grd,save_path=r"H:\test_data\output\meb\gds_test.000")