"""
WRF后处理程序 风 风羽图
20230130 byp
"""
import netCDF4
import wrf
import os
import glob
import re
import math
import matplotlib.pyplot as plt
import matplotlib
import cartopy.crs as ccrs
import cartopy
import numpy as np
from matplotlib.cm import get_cmap
from PIL import Image


# 读取数据###############################################################开始
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)
if os.path.exists(dir_output):
    dir_output = dir_output+'/'+'UV10'
    if not os.path.exists(dir_output):
        os.mkdir(dir_output)
        idx = 0
        while idx < 73:
            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]
            wspd = sd.values[0]

            # 时间
            Start_Date = str(data_nc.START_DATE).replace('_', '\/\/')
            Start_Date_idx = str(uv.coords['Time'].values)[0:19].replace('T', '\/\/')
            Image_Name_idx = str(uv.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)

            # 经纬度范围 直接读取诊断变量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]

            # 底图###################################################################开始
            # 设置中文字体
            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())
            # 设定底图范围
            axes.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
            # 添加标题字符，分左右两边
            title = '风场 10米$\mathbf{(m s-1)}$'  # 图片标题 '海平面气压$\mathbf{(hPa)}$'/'相对湿度 2米$\mathbf{(\%)}$'
            axes.set_title(title, fontsize=5, fontweight='bold', 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{\u00A9HQX}$', 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}
            # 底图###################################################################结束
            matplotlib.pyplot.contourf(wrf.to_np(lon_uv), wrf.to_np(lat_uv), wrf.to_np(wspd), 10,
                                       transform=ccrs.PlateCarree(), cmap=get_cmap("jet"))
            axes.barbs(lon_uv, lat_uv, u, v, length=3, linewidth=0.2,
                       sizes=dict(emptybarb=0), transform=ccrs.PlateCarree())

            fig_name = dir_output + "/" + 'UV10' + "_" + 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’}: 目前只有后端支持

            # 显示图片 logo
            # 加载的水印图片
            im1 = Image.open('./img/logo2 100.png')
            # 底图
            imgName = fig_name
            im2 = Image.open(imgName)
            # 将水印照片与商品图片大小调整
            im1 = im1.resize((100, 100))
            im2.paste(im1, (1530, 1380), im1)
            imgName = imgName
            im2.save(imgName)
            matplotlib.pyplot.close()