"""
WRF后处理程序 slp
20230116 byp
"""
import netCDF4
import wrf
import numpy as np
import matplotlib.pyplot
import cartopy.crs as ccrs
import cartopy.feature
from matplotlib.cm import get_cmap
import os
import glob
import re


# 读取数据###############################################################开始
var = 'SLP' # 海平面气压
dir_name = "./data/WRF/"
dir_output = './output/image/'
list_of_files = glob.glob(dir_name+'*')
latest_file = max(list_of_files, key=os.path.getctime)  # 最新文件
mat = re.search(r"(\d{4}-\d{1,2}-\d{1,2}_\d{1,2})", latest_file)  # 正则匹配时间
dir_time_file = 'wrfout_d01_'+latest_file[mat.span()[0]:]
dir_time = latest_file[mat.span()[0]:mat.span()[1]].replace('-', '').replace('_', '')
data_nc = netCDF4.Dataset(dir_name+dir_time_file)
print(dir_time)
dir_output = dir_output+dir_time
if not os.path.exists(dir_output):
    os.mkdir(dir_output)
    dir_output = dir_output+'/'+var
    os.mkdir(dir_output)
    idx = 0
    while idx < 73:
        slp = wrf.getvar(data_nc, "slp", timeidx=idx)  # 读取海表面气压
        smooth_slp = wrf.smooth2d(slp, 3, cenweight=4)
        lats, lons = wrf.latlon_coords(slp)
        # 时间
        Start_Date = str(data_nc.START_DATE).replace('_', '\/\/')
        Start_Date_idx = str(slp.coords['Time'].values)[0:19].replace('T', '\/\/')
        Image_Name_idx = str(slp.coords['Time'].values)[0:19].replace(':', '_')
        Image_Name_idx = Image_Name_idx.replace('-', '_')
        Image_Name_idx = Image_Name_idx.replace('T', '_')
        print(Image_Name_idx)
        # 经纬度范围
        lat_SIZE = wrf.getvar(data_nc, "XLAT").shape[0]
        lon_SIZE = wrf.getvar(data_nc, "XLONG").shape[0]
        # print(lon_SIZE, lat_SIZE)
        lon_min = wrf.getvar(data_nc, "XLONG")[0][0].values
        lat_min = wrf.getvar(data_nc, "XLAT")[0][0].values
        lon_max = wrf.getvar(data_nc, "XLONG")[0][lon_SIZE - 1].values
        lat_max = wrf.getvar(data_nc, "XLAT")[lat_SIZE - 1][0].values
        # print(lat_min1.values, lon_min1.values, lat_max1.values, lon_max1.values)

        # 底图###################################################################开始
        # 设置中文字体
        matplotlib.pyplot.rcParams['font.sans-serif'] = ['SimHei']

        fig = matplotlib.pyplot.figure(figsize=(4, 4), dpi=150)
        axes = matplotlib.pyplot.subplot(1, 1, 1, projection=ccrs.PlateCarree())
        # 设定底图范围
        # lon_min, lon_max, lat_min, lat_max = 104.1, 135.9, 13.18, 41.04
        axes.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
        # 添加标题字符，分左右两边
        axes.set_title('海平面气压$\mathbf{(hPa)}$', fontsize=5, loc='left', y=0.98)
        axes.set_title('$\mathbf{WRF\/\/V4.0}$', fontsize=4, loc='right', y=0.98)
        # 左下角文字
        axes.text(0.01, -0.02, '$\mathbf{Valid:' + Start_Date_idx + '}$', transform=axes.transAxes,
                  ha='left', va='center', fontsize=4, color='black', fontweight='bold')
        # 右下角文字
        axes.text(.99, -0.02, '$\mathbf{Initial:' + str(Start_Date) + '}$', transform=axes.transAxes,
                  ha='right', va='center', fontsize=4, color='black', fontweight='bold')
        # 右下角文字
        axes.text(.99, 0.02, '$\mathbf{\u00A9BYP}$', transform=axes.transAxes,
                  ha='right', va='center', fontsize=4, color='black', fontweight='bold')
        # 添加地物
        axes.add_feature(cartopy.feature.COASTLINE.with_scale('10m'), lw=0.4)
        # 添加坐标
        gl = axes.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, linestyle=":", linewidth=0.3, x_inline=False,
                            y_inline=False, color='gray')
        gl.top_labels, gl.bottom_labels, gl.right_labels, gl.left_labels = False, False, False, False
        # gl.xformatter = LONGITUDE_FORMATTER  # 使横坐标转化为经纬度格式 无效
        # gl.yformatter = LATITUDE_FORMATTER
        gl.xlocator = matplotlib.ticker.FixedLocator(np.arange(lon_min, lon_max, 2.5))  # 设置坐标轴范围和间隔
        gl.ylocator = matplotlib.ticker.FixedLocator(np.arange(lat_min, lat_max, 2.5))
        gl.xlabel_style = {'size': 4}  # 修改经纬度字体大小
        gl.ylabel_style = {'size': 4}
        # 底图###################################################################结束
        slp_levels = np.arange(980., 1080., 2.5)
        # Get the cartopy mapping object
        cart_proj = wrf.get_cartopy(slp)
        # Make the contour outlines and filled contours for the smoothed sea level pressure.
        line = matplotlib.pyplot.contour(wrf.to_np(lons), wrf.to_np(lats), wrf.to_np(smooth_slp), 10, levels=slp_levels,
                                         colors="black", linewidths=0.2, transform=ccrs.PlateCarree())
        matplotlib.pyplot.contourf(wrf.to_np(lons), wrf.to_np(lats), wrf.to_np(smooth_slp), 10,
                                   transform=ccrs.PlateCarree(), cmap=get_cmap("jet"))
        matplotlib.pyplot.clabel(line, inline=True, fontsize=4)

        fig_name = dir_output + "/SLP_" + Image_Name_idx + ".jpg"
        idx = idx + 1
        matplotlib.pyplot.savefig(fig_name,
                                  dpi=512,  # 分辨率，每英寸的点数 512
                                  bbox_inches='tight',
                                  pad_inches=0.1,  # (默认: 0.1)所保存图形周围的填充量
                                  transparent=True,
                                  facecolor='auto',  # 默认 auto
                                  edgecolor='r',  # 默认 auto
                                  # papertype='letter',  # 3.6之后不支持了。纸张大小，仅支持postscript输出。取值范围为:
                                  # {'letter', 'legal', 'executive', 'ledger', 'a0' - 'a10', 'b0' - 'b10'}。默认值为None。
                                  orientation='portrait')  # {‘landscape,’ ‘portrait’}: 目前只有后端支持
        matplotlib.pyplot.close()