# coding: utf-8
import cinrad
import os
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from matplotlib.colorbar import ColorbarBase
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.io import shapereader
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
from netCDF4 import Dataset
from wrf import (to_np, getvar, smooth2d, get_cartopy, cartopy_xlim,
                 cartopy_ylim, latlon_coords)
from metpy.interpolate import (interpolate_to_grid, remove_nan_observations,
                               remove_repeat_coordinates)
from scipy.interpolate import griddata

def search(s, path=os.path.abspath('.')):
    for z in os.listdir(path):
        if os.path.isdir(path + os.path.sep + z):# os.path.sep
            path2 = os.path.join(path, z) #；os.path.join()
            search(s, path2)
        elif os.path.isfile(path + os.path.sep + z): 
            if s in z:
                filenames.append(os.path.join(path, z))
def plotMapdiff(var,var1,extent=None):
    
    proj = ccrs.PlateCarree(central_longitude=118)
    fig, ax = plt.subplots(figsize=(8.0, 8.0),subplot_kw=dict(projection=proj))
    #plt.subplots_adjust(left=0.05,right=0.95,top=0.90,bottom=0.05,wspace=0.15,hspace=0.05)
    #fig.set_facecolor('w')
    plt.style.use('seaborn-bright')
    lon, lat, var = var.lon, var.lat, var.data
    vara=var.reshape(-1)
    lata=lat.reshape(-1)
    lona=lon.reshape(-1)
    xp0, yp0, _ = proj.transform_points(ccrs.Geodetic(), lona, lata).T
    #print(var11.T.shape)
    gridx0, gridy0, dbz0 = interpolate_to_grid(xp0, yp0,vara, interp_type='linear',hres=0.01)

    lat1, lon1 = latlon_coords(var1)
    lat11=to_np(lat1)
    lon11=to_np(lon1)
    var11=to_np(var1)
    var111=var11.reshape(-1)
    lat111=lat11.reshape(-1)
    lon111=lon11.reshape(-1)
    xp, yp, _ = proj.transform_points(ccrs.Geodetic(), lon111, lat111).T
    gridx1, gridy1, dbz1 = interpolate_to_grid(xp, yp,var111, interp_type='linear',hres=0.01)
    points1=gridx1.reshape(-1)
    points2=gridy1.reshape(-1)
    points=(points1,points2)
    dbz11=dbz1.reshape(-1)
    #Zoom in
    dbz = griddata(points, dbz11, (gridx0,gridy0), method='nearest')
    diff=dbz0-dbz
    diff = np.ma.masked_where(np.isnan(diff), diff)
    
    #ax.background_patch.set_fill(False)
    if not extent:
        x_min, x_max, y_min, y_max = lon.min(), lon.max(), lat.min(), lat.max()
    else:
        x_min, x_max, y_min, y_max = extent[0], extent[1], extent[2], extent[3]
    ax.set_extent([x_min, x_max, y_min, y_max], ccrs.PlateCarree())
    #Add map features
    root = os.path.join("/public/home/hysplit/software/anaconda3/lib/python3.6/site-packages/cinrad-1.2-py3.6.egg/cinrad", 'shapefile')
    flist = [os.path.join(root, i) for i in ['County.shp', 'City.shp', 'Province.shp']]
    shps = [shapereader.Reader(i).geometries() for i in flist]
    line_colors = ['lightgrey', 'grey', 'black']
    ax.add_geometries(shps[0], ccrs.PlateCarree(), edgecolor=line_colors[0], facecolor='None', zorder=1, linewidth=0.5)
    ax.add_geometries(shps[1], ccrs.PlateCarree(), edgecolor=line_colors[1], facecolor='None', zorder=1, linewidth=0.7)
    ax.add_geometries(shps[2], ccrs.PlateCarree(), edgecolor=line_colors[2], facecolor='None', zorder=1, linewidth=1)
    #ax.coastlines(resolution='10m', color=line_colors[2], zorder=1, linewidth=1)
 
    levels = np.arange(-30,35,5)
    print(levels)
    #norm = mpl.colors.Normalize(vmin=0, vmax=70)
    ct=ax.contourf(gridx0, gridy0, diff,levels=levels,  extend='both',cmap="coolwarm")
    #ax.set_title("obs_Reflectivity(dBZ)", {"fontsize" : 20}) 
    #Add lat/lon gridlines every 20° to the map
    gl= ax.gridlines(draw_labels=True,linewidth=2, color='gray', alpha=0.5, linestyle='--') 
    gl.xlabels_top = False
    gl.xlabels_bottom = False
    gl.ylabels_right = False
    gl.ylabels_left = False
    gl.xlocator = mticker.FixedLocator([118,118.5,119,119.5,120])
    gl.ylocator = mticker.FixedLocator([31,31.5,32,32.5,33])
    #'''
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER
    gl.xlabel_style = {'size': 15}#, 'weight': 'bold'}#'color': 'gray'}
    gl.ylabel_style = {'size': 15}
    #cax = fig.add_axes([1.01, 0.75, 0.02, 0.2])
    cb = plt.colorbar(ct,extendfrac=1/10,extendrect=True,shrink=0.6,pad=0.1,orientation='horizontal')#extendrect=True
    
    #cb.ax.tick_params(labelsize=15) 
    #cb.set_ticklabels(levels)
    #plt.clabel(ct,levels, inline=True, fmt='%1i', fontsize=16)
    #fig.text(0.5, -0.04, 'Latitude', ha='center',fontsize=20)
    #fig.text(-0.02, 0.35, 'Longitude', va='center', rotation='vertical',fontsize=20)
    #'''
# Add titles
    #ax.set_title("model_Reflectivity(dBZ)", {"fontsize" : 20})
    
    return fig

def get_model(filename,i=1):
    f = Dataset(filename)
    dbz = getvar(f, "mdbz")
    return dbz

def get_diff(filename,filename1):
    f = Dataset(filename)
    f1 = Dataset(filename1)
    dbz = getvar(f, "mdbz")
    dbz1 = getvar(f1, "mdbz")
    diff = dbz-dbz1
    return diff
def get_obs(filename):
    f = cinrad.io.CinradReader(filename)
    rl = [f.get_data(i, 200, 'REF') for i in f.angleindex_r]
    cr = cinrad.easycalc.quick_cr(rl) #计算组合反射率
    cr.name = 'Nanjing'
    print(cr)
    return cr

filenames=[]
search('case_d04', '.')
filenames=sorted(filenames)#;防止顺序不对
filenames0=filenames[0:3]
print(filenames0)
filenames1=filenames[3:6]
print(filenames1)
filenames2=filenames[6:9]
print(filenames2)
filenames3=filenames[9:12]
print(filenames3)
filename_obs=['Z_RADR_I_Z9250_20170609230400_O_DOR_SA_CAP.bin','Z_RADR_I_Z9250_20170610045300_O_DOR_SA_CAP.bin','Z_RADR_I_Z9250_20170610080400_O_DOR_SA_CAP.bin']
crsobs=[get_obs(filename) for filename in filename_obs]
crsmod=[get_model(filename) for filename in filenames0]
NUM=['P','Q','R']
for i in range(0,3):
    plotMapdiff(crsobs[i],crsmod[i],extent=[118,119.55,31.25,32.7])
    plt.text(118.05,32.60,NUM[i],color='k',fontsize=30,
         horizontalalignment='left',
         transform=ccrs.Geodetic())
    #plt.show()
    plt.savefig(NUM[i]+"b.png",dpi=600,bbox_inches="tight")
crsmod=[get_model(filename) for filename in filenames1]
NUM=['S','T','U']
for i in range(0,3):

    plotMapdiff(crsobs[i],crsmod[i],extent=[118,119.55,31.25,32.7])
    plt.text(118.05,32.60,NUM[i],color='k',fontsize=30,
         horizontalalignment='left',
         transform=ccrs.Geodetic())
    #plt.show()
    plt.savefig(NUM[i]+"a.png",dpi=600,bbox_inches="tight")
crsmod=[get_model(filename) for filename in filenames2]
NUM=['V','W','X']
for i in range(0,3):

    plotMapdiff(crsobs[i],crsmod[i],extent=[118,119.55,31.25,32.7])
    plt.text(118.05,32.60,NUM[i],color='k',fontsize=30,
         horizontalalignment='left',
         transform=ccrs.Geodetic())
    #plt.show()
    plt.savefig(NUM[i]+"a.png",dpi=600,bbox_inches="tight")

crsmod=[get_model(filename) for filename in filenames3]
NUM=['Y','Z','&']
for i in range(0,3):

    plotMapdiff(crsobs[i],crsmod[i],extent=[118,119.55,31.25,32.7])
    plt.text(118.05,32.60,NUM[i],color='k',fontsize=30,
         horizontalalignment='left',
         transform=ccrs.Geodetic())
    #plt.show()
    plt.savefig(NUM[i]+"a.png",dpi=600,bbox_inches="tight")

diff