"""
WRF后处理程序 风变量对比
20230130 byp
"""
import netCDF4
import wrf
import os
import glob
import re
import math
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy


# 读取数据###############################################################开始
dir_name = "./data/WRF/"
dir_output = './output/txt/'
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)

idx = 0
i = 0

uv = wrf.getvar(data_nc, 'uvmet10', timeidx=idx)
sd = wrf.getvar(data_nc, 'wspd_wdir10', timeidx=idx)
sd10 = wrf.getvar(data_nc, 'uvmet10_wspd_wdir', timeidx=idx)
u = uv.values[0]
v = uv.values[1]

# 经纬度范围 直接读取诊断变量XLAT和XLONG
lon_uv = uv.coords['XLONG'].values
lat_uv = uv.coords['XLAT'].values
lat_SIZE = lat_uv.shape[0]
lon_SIZE = lon_uv.shape[0]
lon_min = lon_uv[0][0]
lon_max = lon_uv[0][lon_SIZE-1]
lat_min = lat_uv[0][0]
lat_max = lat_uv[lat_SIZE-1][0]
print(str(lat_min)+' '+str(lon_min)+' '+str(lat_max)+' '+str(lon_max))

fig = plt.figure(figsize=(10, 5))
axes = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
# 设定底图范围
axes.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
axes.stock_img()
# 添加地物
axes.add_feature(cartopy.feature.COASTLINE.with_scale('10m'), lw=0.4)
axes.barbs(lon_uv, lat_uv, u, v, length=3,
           sizes=dict(emptybarb=0.25, spacing=0.2, height=0.5),
           linewidth=0.95, transform=ccrs.PlateCarree())
plt.show()