# plot landuse and terrain

from netCDF4 import Dataset
import numpy as np
from matplotlib.colors import BoundaryNorm
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.colors import LinearSegmentedColormap
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.io.shapereader import Reader
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
from wrf import (getvar, to_np, latlon_coords, get_cartopy,
                 ALL_TIMES)
from mypath import lambertticks as lt
from matplotlib.patches import Rectangle
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter, AutoMinorLocator)
import matplotlib.lines as lines

reader = Reader("D:/Dhy/MAP/China_provinces/China_provinces")
reader_w = Reader("D:/Dhy/MAP/world_countries/world_countries")
provinces = cfeature.ShapelyFeature(reader.geometries(), ccrs.PlateCarree(), edgecolor='k', facecolor='none')
world = cfeature.ShapelyFeature(reader_w.geometries(), ccrs.PlateCarree(), edgecolor='k', facecolor='none')
kzmin = ['0.2', '0.5']
days = ['08-10','11-16','20-21','22-24']
day = ['07', '08', '09', '10', '11', '12', '13', '14', '19', '20', '21', '22', '23']
pbl = ['AC', 'BL']
pathf1 = 'F:/wrfout/xiangmu/wrfoutACORG/'
pathf2 = 'F:/wrfout/xiangmu/wrfoutkf/'
d02path= 'F:/wrfout/wrfchem/wrfout0120/wrfout_d02_2014-01-20_000000'
d01path= 'F:/wrfout/wrf_bj/run1-acm2/wrfout_d01_2014-01-20_000000'
d03path= 'F:/wrfout/xiangmu/wrfoutkf/wrfout_d03_2014-01-20_00_00_00'


def setticks(lat,lon,ax):
    # Set the map bounds and plot ticks
    print([to_np(lat).max(), to_np(lon).max(), to_np(lat).min(), to_np(lon).min()])
    minlon=to_np(lon).min()
    maxlon=to_np(lon).max()
    minlat=to_np(lat).min()
    maxlat=to_np(lat).max()
    ax.set_extent([lon[0,0], lon[0,-1], lat[0,-1], lat[-1,0]], crs=ccrs.PlateCarree())
    fig.canvas.draw()
    xticks = list(range(minlon.astype(int)-1, maxlon.astype(int)+3, 1))
    yticks = list(range(minlat.astype(int)-1, maxlat.astype(int)+3, 1))
    ax.gridlines(xlocs=xticks, ylocs=yticks)
    ax.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
    ax.yaxis.set_major_formatter(LATITUDE_FORMATTER)
    lt.lambert_xticks(ax, xticks)
    lt.lambert_yticks(ax, yticks)
    ax.tick_params(labelsize=15)

def setticks_domain1(lat,lon,ax):
    # Set the map bounds and plot ticks
    print([to_np(lat).max(), to_np(lon).max(), to_np(lat).min(), to_np(lon).min()])
    minlon=to_np(lon).min()-2
    maxlon=to_np(lon).max()+2
    minlat=to_np(lat).min()-2
    maxlat=to_np(lat).max()+2
    ax.set_extent([minlon, maxlon, minlat, maxlat], crs=ccrs.PlateCarree())
    fig.canvas.draw()
    xticks = list(range(minlon.astype(int)-6, maxlon.astype(int)+9, 5))
    yticks = list(range(minlat.astype(int)-6, maxlat.astype(int)+9, 5))
    ax.gridlines(xlocs=xticks, ylocs=yticks)
    ax.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
    ax.yaxis.set_major_formatter(LATITUDE_FORMATTER)
    lt.lambert_xticks(ax, xticks)
    lt.lambert_yticks(ax, yticks)
    ax.tick_params(labelsize=15)

def readvar2d(pathf1, var2d):
    ncfile = Dataset(pathf1)
    t2d = getvar(ncfile, var2d)
    lat = getvar(ncfile, "XLAT")
    lon = getvar(ncfile, "XLONG")
    return t2d, lat, lon



# get var
lu_index, lat, lon = readvar2d(d03path, 'LU_INDEX')
lu_index2d = lu_index[:, :]

hgt01, lat01, lon01 = readvar2d(d01path, 'HGT')
hgt02, lat02, lon02 = readvar2d(d02path, 'HGT')
hgt03, lat03, lon03 = readvar2d(d03path, 'HGT')


###### Plot Landuse
# Get the map projection information
cart_proj = get_cartopy(lon)

# define Landuse Colors
colors = ['springgreen', 'mediumseagreen', 'yellowgreen', 'yellow',         'greenyellow',
          'gold',        'darkkhaki',      'darkorange',  'sandybrown',     'mediumpurple',
          'darkcyan',    'aqua',           'red',         'fuchsia',        'grey',
          'brown',       'dodgerblue',     'pink',        'lightsteelblue', 'indigo']
cm = LinearSegmentedColormap.from_list('mylist', colors, 20)

# plot landusef
fig = plt.figure(figsize=(7,7))
ax = plt.axes(projection=cart_proj)

# add the states and coastlines
ax.add_feature(provinces, linewidth=0.5, edgecolor="black")

levels = np.arange(0.5, 21.5, 1)
cmap = plt.get_cmap(cm)
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)
wspd_contours = ax.pcolormesh(to_np(lon), to_np(lat), to_np(lu_index2d),
                              cmap=get_cmap(cm), norm=norm, #get_cmap("YlOrRd"),get_cmap("terrain")
                              transform=ccrs.PlateCarree())
# cbar = plt.colorbar(wspd_contours, ax=ax, orientation="vertical", pad=.05)
setticks(lat, lon, ax)
fig.subplots_adjust(top=0.985,bottom=0.045,left=0.013,right=0.988)


def plot_colortablea(colors, title, sort_colors=True, emptycols=0):

    cell_width = 212
    cell_height = 22
    swatch_width = 48
    margin = 12
    topmargin = 40

    # Sort colors by hue, saturation, value and name.

    names = colors
    names = ['Evergreen Needleleaf Forest', 'Evergreen Broadleaf Forest',
             'Deciduous Needleleaf Forest', 'Deciduous Broadleaf Forest',
             'Mixed Forest',                'Closed Shrublands',
             'Open Shrublands',             'Woody Savannas',
             'Savannas',                    'Grasslands',
             'Permanent Wetlands',          'Croplands',
             'Urban and Built-Up',          'Cropland/Natural Vegetation Mosaic',
             'Snow and Ice',                'Barren or Sparsely Vegetated',
             'Water',                       'Wooded Tundra',
             'Mixed Tundra',                'Barren Tundra']

    n = len(names)
    ncols = 1 - emptycols
    nrows = n // ncols + int(n % ncols > 0)

    width = cell_width * 4 + 2 * margin
    height = cell_height * nrows + margin + topmargin
    dpi = 72

    fig, ax = plt.subplots(figsize=(width / dpi, height / dpi), dpi=dpi)
    fig.subplots_adjust(margin/width, margin/height,
                        (width-margin)/width, (height-topmargin)/height)
    ax.set_xlim(0, cell_width * 4)
    ax.set_ylim(cell_height * (nrows-0.5), -cell_height/2.)
    ax.yaxis.set_visible(False)
    ax.xaxis.set_visible(False)
    ax.set_axis_off()
    #ax.set_title(title, fontsize=24, loc="left", pad=10)

    for i, name in enumerate(names):
        row = i % nrows
        col = i // nrows
        y = row * cell_height

        swatch_start_x = cell_width * col
        text_pos_x = cell_width * col + swatch_width + 7

        ax.text(text_pos_x, y, name, fontsize=14,
                horizontalalignment='left',
                verticalalignment='center')

        ax.add_patch(
            Rectangle(xy=(swatch_start_x, y-9), width=swatch_width,
                      height=18, facecolor=colors[i])
        )

    return fig

plot_colortablea(colors, "Colors")



##### Plot terrain
fig = plt.figure(figsize=(17,7))
ax1 = fig.add_subplot(1,2,1,projection=cart_proj)
ax2 = fig.add_subplot(1,2,2,projection=cart_proj)

fig.subplots_adjust(top=0.985,bottom=0.045,left=0.038,right=0.988)

# add the states and coastlines
ax1.add_feature(provinces, linewidth=0.5, edgecolor="black")
ax1.add_feature(world, linewidth=0.5, edgecolor="black")
ax2.add_feature(provinces, linewidth=0.5, edgecolor="black")

levels = [50, 100, 200, 400, 800, 1200, 1600, 2000]
#levels = np.arange(0,2001,1)

cmap = plt.get_cmap('terrain')
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)
wspd_contours = ax1.pcolormesh(to_np(lon01), to_np(lat01), to_np(hgt01[:, :]),#to_np(hgt),#to_np(pm25[0, 0, :, :]),
                            #vmin=0,vmax=0.6, #levels=levels,
                            cmap=get_cmap("terrain"),norm=norm, #get_cmap("YlOrRd"),get_cmap("terrain")
                            transform=ccrs.PlateCarree())

setticks_domain1(lat01, lon01, ax1)
# plt.title("T2diff(K)", fontsize=15)



# boudary
def boundary(data):
    bd = np.array(data)
    bd[:,:] = 0
    bd[0,:] = 1
    bd[-1,:] = 1
    bd[:,0] = 1
    bd[:,-1] = 1
    return np.where(bd==1)

location1 = boundary(lat01)
location2 = boundary(lat02)
location3 = boundary(lat03)


loc = [[0,0,-1,-1,0],[0,-1,-1,0,0]]
loc = np.array(loc)
ax1.plot(to_np(lon01)[loc[0],loc[1]],to_np(lat01)[loc[0],loc[1]],transform=ccrs.PlateCarree(),color='k', linewidth=2)
ax1.plot(to_np(lon02)[loc[0],loc[1]],to_np(lat02)[loc[0],loc[1]],transform=ccrs.PlateCarree(),color='k', linewidth=2)
ax1.plot(to_np(lon03)[loc[0],loc[1]],to_np(lat03)[loc[0],loc[1]],transform=ccrs.PlateCarree(),color='k', linewidth=2)

ax1.text(115.15, 40.63, 'D03', fontsize=20, color='r', fontdict={'weight': 400}, transform=ccrs.PlateCarree())
ax1.text(113.65, 42.08, 'D02', fontsize=20, color='r', fontdict={'weight': 400},transform=ccrs.PlateCarree())
ax1.text(110.30, 44.17, 'D01', fontsize=20, color='r', fontdict={'weight': 400},transform=ccrs.PlateCarree())

#axs[1] = plt.axes(projection=cart_proj)
wspd_contours3 = ax2.pcolormesh(to_np(lon03), to_np(lat03), to_np(hgt03[:, :]),#to_np(hgt),#to_np(pm25[0, 0, :, :]),
                            #vmin=0,vmax=0.6, #levels=levels,
                            cmap=get_cmap("terrain"),norm=norm, #get_cmap("YlOrRd"),get_cmap("terrain")
                            transform=ccrs.PlateCarree())
setticks(lat03, lon03, ax2)
fig.add_artist(lines.Line2D([0.213, 0.580], [0.576, 0.986],color='r',linewidth=1))
fig.add_artist(lines.Line2D([0.213, 0.580], [0.413, 0.045],color='r',linewidth=1))
#fig.add_artist(lines.Line2D([0, 1], [0, 1],color='k',linewidth=2))
'''''
x1, x2, y1, y2 = -1.631e5, 9.4e4, -1.21e5, 1.3e5
axins.set_xlim(x1, x2)
axins.set_ylim(y1, y2)
ax.indicate_inset_zoom(axins, edgecolor="black")
'''
# color bar plot
fig, ax = plt.subplots(figsize=(12,10))
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)
fig.subplots_adjust(top=0.985)
fig.subplots_adjust(bottom=0.045)
fig.subplots_adjust(left=0.020)
fig.subplots_adjust(right=0.950)
ax.set_xticks([])
ax.set_yticks([])
cbar = plt.colorbar(wspd_contours, ax=ax, orientation="horizontal")
cbar.ax.set_xlabel('terrain (m)', fontsize=25)
cbar.set_ticks(levels)
cbar.ax.tick_params(labelsize=25) #设置colorbar字体大小

plt.show()


