"""
WRF后处理程序 slp
20230116 byp
"""
import netCDF4
import wrf
import os
import glob
import re
import matplotlib
import matplotlib.pyplot
import cartopy
from wrf import (getvar, interplevel, to_np, latlon_coords, get_cartopy,
                 cartopy_xlim, cartopy_ylim)


# 读取数据###############################################################开始
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)  # 最新文件
print(latest_file)
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)

# 经纬度范围
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


u = wrf.getvar(data_nc, 'U10', timeidx=1)
# u = u.values[1][10]
print(u.coords.values())
v = wrf.getvar(data_nc, 'V10', timeidx=1)

# Get the lat/lon coordinates
lats, lons = wrf.latlon_coords(u)

fig = matplotlib.pyplot.figure(figsize=(8, 8), dpi=150)
axes = matplotlib.pyplot.subplot(1, 1, 1, projection=cartopy.crs.PlateCarree())
axes.set_extent([lon_min, lon_max, lat_min, lat_max], crs=cartopy.crs.PlateCarree())
axes.add_feature(cartopy.feature.COASTLINE.with_scale('10m'), lw=0.4)

axes.barbs(lons, lats, u, v, length=3.8, linewidth=0.55)

matplotlib.pyplot.show()