import netCDF4
import wrf
import matplotlib.pyplot as plt
import matplotlib.ticker
import numpy as np
import cartopy.crs as ccrs
import cartopy.feature as cf
from matplotlib.cm import get_cmap


# 读取数据###############################################################
data_nc = netCDF4.Dataset(r"./data/WRF/wrfout_d02_2023-01-15_12%3A00%3A00")
slp = wrf.getvar(data_nc, "slp")  # 读取计算海表面气压
smooth_slp = wrf.smooth2d(slp, 3, cenweight=4)
time_slp = str(slp.coords['Time'].values)[0:19].replace('T', '\/\/')
lats, lons = wrf.latlon_coords(slp)
# u, v = wrf.uvmet(data_nc)

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体

fig = plt.figure(figsize=(4, 4), dpi=150)
axes = plt.subplot(1, 1, 1, projection=ccrs.PlateCarree())

lon_min, lon_max, lat_min, lat_max = 103, 127.5, 12, 28
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:2022-12-17\/\/12:00:00}$', transform=axes.transAxes,
          ha='left', va='center', fontsize=4, color='black', fontweight='bold')
# 右下角文字
axes.text(.99, -0.02, '$\mathbf{Initial:'+str(time_slp)+'}$', 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(cf.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}
################################################################
# Get the cartopy mapping object
cart_proj = wrf.get_cartopy(slp)
# Make the contour outlines and filled contours for the smoothed sea level pressure.
plt.contour(wrf.to_np(lons), wrf.to_np(lats), wrf.to_np(smooth_slp), 10, colors="black", linewidths=0.2,
            transform=ccrs.PlateCarree())
plt.contourf(wrf.to_np(lons), wrf.to_np(lats), wrf.to_np(smooth_slp), 10,
             transform=ccrs.PlateCarree(),
             cmap=get_cmap("jet"))

fig_name = "./test_WRF.jpg"
plt.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’}: 目前只有后端支持