#!/usr/bin/env python
# coding: utf-8

# # This notebook exhibits a few plots for the data tiers

# In[2]:


#This notebook was created on local systems.
#To use this notebook on SageMaker, users will
#need to load the data from the wrf-cmip6-noversioning
#bucket using the specific AWS library calls. For more information,
#see https://docs.aws.amazon.com/sagemaker/latest/dg/gs-setup-working-env.html


# In[32]:


#Functions and otherwise...

from netCDF4 import Dataset
import pylab as P
import numpy as np
from numpy import ma
from numpy import dtype
import matplotlib
import cartopy.crs as crs
import cartopy.feature as cfeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import matplotlib.colors as colors
import matplotlib.ticker as mticker
import matplotlib.gridspec as gridspec
import xarray as xr
import pandas as pd
import os
from scipy.interpolate import griddata
import wrf
from wrf import (getvar, to_np, vertcross, smooth2d, CoordPair, GeoBounds,
                 get_cartopy, latlon_coords, cartopy_xlim, cartopy_ylim)
import xesmf as xe
import dask
import gc
import datetime
import time
from scipy import stats
import regionmask

states = cfeature.NaturalEarthFeature(category='cultural', scale='50m',
                                      facecolor='none',
                                      name='admin_1_states_provinces_shp')

#Read metadata (wrfinput file)
def _metaread(dir_meta,domain):
 file = "%swrfinput_%s" %(dir_meta,domain)
 data = xr.open_dataset(file)
 lat = data.variables["XLAT"]
 lon = data.variables["XLONG"]
 z = data.variables["HGT"]

 return (lat,lon,z,file)

#WRF reader function
def _wrfread(prefix,dir,var,domain):

 all_files = sorted(os.listdir(dir))

 anal_files = []
 for ii in all_files:
  if ii.startswith(var+"."):
   if domain in ii:
    if prefix in ii:
     anal_files.append(dir+str(ii))

 del all_files

 nf  = len(anal_files)

 data = xr.open_mfdataset(anal_files, combine="by_coords")
 var_read = data.variables[var]
 day = data.variables["day"].values
 nt = len(day)

 day1 = pd.to_datetime(str(int(day[0])), format='%Y-%m-%d')
 day2 = pd.to_datetime(str(int(day[nt-1])), format='%Y-%m-%d')
 dates = pd.date_range(day1,day2,freq="D")

 #Mask array setting leap years = True
 is_leap_day = (dates.month == 2) & (dates.day == 29)
 dates = dates[~is_leap_day]

 var_read = xr.DataArray(var_read, dims=['day','lat2d','lon2d'])
 var_read['day'] = dates    #year doesn't matter here

 return (var_read)

#WRF tier 3 reader function
def _wrfread_gcm(model,gcm,variant,dir,var,domain):

 all_files = sorted(os.listdir(dir))

 anal_files = []
 for ii in all_files:
  if ii.startswith(var+".") and model in ii and gcm in ii \
        and variant in ii and domain in ii:
   if domain in ii:
     anal_files.append(dir+str(ii))

 del all_files

 nf  = len(anal_files)

 data = xr.open_mfdataset(anal_files, combine="by_coords")
 var_read = data.variables[var]
 day = data.variables["day"].values
 nt = len(day)

 day1 = pd.to_datetime(str(int(day[0])), format='%Y-%m-%d')
 day2 = pd.to_datetime(str(int(day[nt-1])), format='%Y-%m-%d')
 dates = pd.date_range(day1,day2,freq="D")

 #Mask array setting leap years = True
 ''' is_leap_day = (dates.month == 2) & (dates.day == 29)
 dates = dates[~is_leap_day]'''

 var_read = xr.DataArray(var_read, dims=['day','lat2d','lon2d'])
 var_read['day'] = dates    #year doesn't matter here

 return (var_read)

def screen_times(data,date_start,date_end):
    #Dimensions should be "time"
    dask.config.set(**{'array.slicing.split_large_chunks': True})
    data = data.sel(time=~((data.time.dt.month < date_start[1]) & (data.time.dt.year <= date_start[0])))
    data = data.sel(time=~((data.time.dt.year < date_start[0])))
    data = data.sel(time=~((data.time.dt.month >= date_end[1]) & (data.time.dt.year >= date_end[0])))
    data = data.sel(time=~((data.time.dt.year > date_end[0])))

    return(data)

def screen_times_wrf(data,date_start,date_end):
    #Dimensions should be "day"
    dask.config.set(**{'array.slicing.split_large_chunks': True})
    data = data.sel(day=~((data.day.dt.month < date_start[1]) & (data.day.dt.year <= date_start[0])))
    data = data.sel(day=~((data.day.dt.year < date_start[0])))
    data = data.sel(day=~((data.day.dt.month >= date_end[1]) & (data.day.dt.year >= date_end[0])))
    data = data.sel(day=~((data.day.dt.year > date_end[0])))
    
    return(data)

def screen_times_wrf_year_only(data,date_start,date_end):
    #Dimensions should be "day"
    dask.config.set(**{'array.slicing.split_large_chunks': True})
    data = data.sel(day=~((data.day.dt.year < date_start[0])))
    data = data.sel(day=~((data.day.dt.year > date_end[0])))
    
    return(data)

print ("Functions loaded")


# # Read in the metadata from the wrfinput files

# In[6]:


get_ipython().run_cell_magic('time', '', '\n#These files contain all static fields\ndomain = "d01"\ndir_meta = "/mnt/parm/raid/data/WRF_2020/meta/meta_new/"\nlat1, lon1, z1, file = _metaread(dir_meta,domain)\nlon_wrf = lon1[0,:,:]\nlat_wrf = lat1[0,:,:]\nz_wrf = z1[0,:,:]\nlat_wrf = xr.DataArray(lat_wrf, dims=["lat2d", "lon2d"])\nlon_wrf = xr.DataArray(lon_wrf, dims=["lat2d", "lon2d"])\nz_wrf = xr.DataArray(z_wrf, dims=["lat2d", "lon2d"])\n\nlat_wrf_d01, lon_wrf_d01 = lat_wrf, lon_wrf\n\ndomain = "d02"\ndir_meta = "/mnt/parm/raid/data/WRF_2020/meta/meta_new/"\nlat1, lon1, z1, file = _metaread(dir_meta,domain)\nlon_wrf = lon1[0,:,:]\nlat_wrf = lat1[0,:,:]\nz_wrf = z1[0,:,:]\nlat_wrf = xr.DataArray(lat_wrf, dims=["lat2d", "lon2d"])\nlon_wrf = xr.DataArray(lon_wrf, dims=["lat2d", "lon2d"])\nz_wrf = xr.DataArray(z_wrf, dims=["lat2d", "lon2d"])\n\nlat_wrf_d02, lon_wrf_d02 = lat_wrf, lon_wrf\n\ndomain = "d03"\ndir_meta = "/mnt/parm/raid/data/WRF_2020/meta/meta_new/"\nlat1, lon1, z1, file = _metaread(dir_meta,domain)\nlon_wrf = lon1[0,:,:]\nlat_wrf = lat1[0,:,:]\nz_wrf = z1[0,:,:]\nlat_wrf = xr.DataArray(lat_wrf, dims=["lat2d", "lon2d"])\nlon_wrf = xr.DataArray(lon_wrf, dims=["lat2d", "lon2d"])\nz_wrf = xr.DataArray(z_wrf, dims=["lat2d", "lon2d"])\n\nlat_wrf_d03, lon_wrf_d03 = lat_wrf, lon_wrf\n\ndomain = "d04"\ndir_meta = "/mnt/parm/raid/data/WRF_2020/meta/meta_new/"\nlat1, lon1, z1, file = _metaread(dir_meta,domain)\nlon_wrf = lon1[0,:,:]\nlat_wrf = lat1[0,:,:]\nz_wrf = z1[0,:,:]\nlat_wrf = xr.DataArray(lat_wrf, dims=["lat2d", "lon2d"])\nlon_wrf = xr.DataArray(lon_wrf, dims=["lat2d", "lon2d"])\nz_wrf = xr.DataArray(z_wrf, dims=["lat2d", "lon2d"])\n\nlat_wrf_d04, lon_wrf_d04 = lat_wrf, lon_wrf\n\nprint("metadata read for each grid")\n')


# # Tier 1 (*wrfout*) and Tier 2 (*auxhist*) files

# In[4]:


#Visualizing Tier 1 and Tier 2 data is a similar process.
#We thus focus on a given set of hourly (Tier 2) files from era5.


# In[29]:


#For Tier 2

#Read files
dir = "/mnt/parm/raid/data/WRF_2020/era/hourly/2000/d03/"

all_files = sorted(os.listdir(dir))

analysis_files = []
for ii in all_files[0:24]:         #Only focus on a single day for example
 analysis_files.append(dir+str(ii))

del all_files

var = 'T2'

lump_in_time = []

for ifile in analysis_files:
 data = xr.open_dataset(ifile)
 lump_in_time.append(data)

#Finalized merged Dataset from which
#we extract the temperature variable
lump_in_time = xr.concat(lump_in_time, dim='Time')

#Averaged in Time
surface_temperature = lump_in_time.variables[var].mean(dim="Time")

#Now plot

#Custom colorbar
levs =  np.array(np.arange(285,305,2))
mycmap = ["darkblue","blue","aqua","lightyellow",
        "red","purple"]

colors_to_rgb = []
for ii in mycmap:
 aa = colors.to_rgba(ii)
 colors_to_rgb.append(aa)
cmap = LinearSegmentedColormap.from_list("t2",colors_to_rgb,70)

#Define the projection
cart_proj = crs.PlateCarree()

fig = P.figure(figsize=(12,12),dpi=250)

ax1 = P.subplot(1,1,1,projection=cart_proj)

txt = "2 m Temperature"
ax1.set_title(txt,fontsize=20,loc="left")

cs = ax1.pcolormesh(lon_wrf_d03,lat_wrf_d03,surface_temperature[:-1,:-1], \
                cmap=cmap, \
                vmin=np.min(levs),vmax=np.max(levs),shading='flat')

ax = [ax1]

lons = np.arange(-170,-60,4)
lats = np.arange(10,70,4)

for iax in ax:
    iax.coastlines('50m', linewidth=0.8)
    iax.add_feature(states, linewidth=.5, edgecolor="black")
    gl = iax.gridlines(crs=cart_proj,color="black", linestyle="dotted", \
                   x_inline=True,y_inline=True)
    gl.xlocator = mticker.FixedLocator(lons)
    gl.ylocator = mticker.FixedLocator(lats)
    iax.set_xticks(np.arange(-170,-60,4),crs=crs.PlateCarree())
    iax.set_yticks(np.arange(10,70,4),crs=crs.PlateCarree())
    lon_formatter = LongitudeFormatter()
    lat_formatter = LatitudeFormatter()
    iax.xaxis.set_major_formatter(lon_formatter)
    iax.yaxis.set_major_formatter(lat_formatter)
    
    for tick in iax.xaxis.get_major_ticks():
                tick.label.set_fontsize(10)
    for tick in iax.yaxis.get_major_ticks():
                tick.label.set_fontsize(10)

    iax.set_extent([-125,-114,32,43])

#Get positions for cbs
ax1_pos = ax1.get_position()

cbar_ax = fig.add_axes([ax1_pos.x1+0.02, ax1_pos.y0+0.01, 0.02, ax1_pos.y1 - ax1_pos.y0-0.01]) # bottom left corner x,y, and width, height
cb = P.colorbar(cs,cax=cbar_ax,extend="both",orientation="vertical")
cb.ax.tick_params(labelsize=15)


# # Now let's make this more interesting.... use the Tier 3 files to examine the climate change signal for snow water equivalent focusing on a single GCM

# In[54]:


get_ipython().run_cell_magic('time', '', '\ndomain = "d01"\n\nvar = \'snow\'\ngcm = \'mpi-esm1-2-lr\'\nvariant = \'r8i1p1f1\'\n\ndomain = "d01"\nssp = \'ssp370\'\ndate_start_pd, date_end_pd = [1980,9,1], [2010,9,1]\ndate_start_ec, date_end_ec = [2070,9,1], [2100,9,1]\n\nmodel = \'hist\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_pd,date_end_pd)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny1 = var_wrf\npd_array = var_wrf.mean(dim="day")\n\nmodel = \'ssp370\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_ec,date_end_ec)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny2 = var_wrf\nec_array = var_wrf.mean(dim="day")\n\ndiff_d01 = y2 - y1\ndiff_mean_d01 = ec_array - pd_array\n\n#Compute t test\ny1 = np.array(y1)\ny2 = np.array(y2)\n\n#Mask where < 0.1, 0.05, 0.01, computing t- and p-values\nthresh = 0.1\nt, p  = stats.ttest_ind(y2,y1,equal_var=False)\np = np.ma.masked_invalid(p)\n\np1_d01 = np.ma.masked_where(p > thresh, p)\nthresh = 0.05\np2_d01 = np.ma.masked_where(p > thresh, p1_d01)\nthresh = 0.01\np3_d01 = np.ma.masked_where(p > thresh, p2_d01)\n\ndomain = "d02"\nssp = \'ssp370\'\ndate_start_pd, date_end_pd = [1980,9,1], [2010,9,1]\ndate_start_ec, date_end_ec = [2070,9,1], [2100,9,1]\n\nmodel = \'hist\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_pd,date_end_pd)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny1 = var_wrf\npd_array = var_wrf.mean(dim="day")\n\nmodel = \'ssp370\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_ec,date_end_ec)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny2 = var_wrf\nec_array = var_wrf.mean(dim="day")\n\ndiff_d02 = y2 - y1\ndiff_mean_d02 = ec_array - pd_array\n\n#Compute t test\ny1 = np.array(y1)\ny2 = np.array(y2)\n\n#Mask where < 0.1, 0.05, 0.01, computing t- and p-values\nthresh = 0.1\nt, p  = stats.ttest_ind(y2,y1,equal_var=False)\np = np.ma.masked_invalid(p)\n\np1_d02 = np.ma.masked_where(p > thresh, p)\nthresh = 0.05\np2_d02 = np.ma.masked_where(p > thresh, p1_d02)\nthresh = 0.01\np3_d02 = np.ma.masked_where(p > thresh, p2_d02)\n\ndomain = "d03"\nssp = \'ssp370\'\ndate_start_pd, date_end_pd = [1980,9,1], [2010,9,1]\ndate_start_ec, date_end_ec = [2070,9,1], [2100,9,1]\n\nmodel = \'hist\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_pd,date_end_pd)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny1 = var_wrf\npd_array = var_wrf.mean(dim="day")\n\nmodel = \'ssp370\'\ndir = "/mnt/parm/raid/data/WRF_2020/mpi_r8/postprocess/" + domain + "/"\nvar_wrf = _wrfread_gcm(model,gcm,variant,dir,var,domain)\nvar_wrf = screen_times_wrf(var_wrf,date_start_ec,date_end_ec)\nvar_wrf = var_wrf.sel(day=( (var_wrf.day.dt.month == 4) & (var_wrf.day.dt.day == 1) )  )\ny2 = var_wrf\nec_array = var_wrf.mean(dim="day")\n\ndiff_d03 = y2 - y1\ndiff_mean_d03 = ec_array - pd_array\n\n#Compute t test\ny1 = np.array(y1)\ny2 = np.array(y2)\n\n#Mask where < 0.1, 0.05, 0.01, computing t- and p-values\nthresh = 0.1\nt, p  = stats.ttest_ind(y2,y1,equal_var=False)\np = np.ma.masked_invalid(p)\n\np1_d03 = np.ma.masked_where(p > thresh, p)\nthresh = 0.05\np2_d03 = np.ma.masked_where(p > thresh, p1_d03)\nthresh = 0.01\np3_d03 = np.ma.masked_where(p > thresh, p2_d03)\n\n#Plot\nlevs = [-500,-400,-300,-200,-100,-50,-25,-10,0,10,25,50,100,200,300,400,500]\nlevs = np.array(levs)\nmycmap = ["darkblue","blue","aquamarine","white","lightyellow","red","darkred"]\nnorm = matplotlib.colors.BoundaryNorm(levs,len(levs))\n\ncolors_to_rgb = []\nfor ii in mycmap:\n aa = colors.to_rgba(ii)\n colors_to_rgb.append(aa)\ncmap = LinearSegmentedColormap.from_list("t2",colors_to_rgb,len(levs))\n\n#Define the projection\ncart_proj = crs.PlateCarree()\n\nfig = P.figure(figsize=(12,12),dpi=250)\n\nax1 = P.subplot(1,3,1,projection=cart_proj)\n\ntxt = "45-km"\nax1.set_title(txt,fontsize=20,loc="left")\n\nax1.pcolormesh(lon_wrf_d01,lat_wrf_d01,diff_mean_d01[:-1,:-1], \\\n                cmap=cmap,norm=norm, \\\n                shading=\'flat\')\n\nax2 = P.subplot(1,3,2,projection=cart_proj)\n\ntxt = "9-km"\nax2.set_title(txt,fontsize=20,loc="left")\n\nax2.pcolormesh(lon_wrf_d02,lat_wrf_d02,diff_mean_d02[:-1,:-1], \\\n                cmap=cmap,norm=norm, \\\n                shading=\'flat\')\n\nax3 = P.subplot(1,3,3,projection=cart_proj)\n\ntxt = "3-km"\nax3.set_title(txt,fontsize=20,loc="left")\n\ncs = ax3.pcolormesh(lon_wrf_d03,lat_wrf_d03,diff_mean_d03[:-1,:-1], \\\n                cmap=cmap,norm=norm, \\\n                shading=\'flat\')\n\nax = [ax1,ax2,ax3]\n\nlons = np.arange(-170,-60,4)\nlats = np.arange(10,70,4)\n\nfor iax in ax:\n    iax.coastlines(\'50m\', linewidth=0.8)\n    iax.add_feature(states, linewidth=.5, edgecolor="black")\n    gl = iax.gridlines(crs=cart_proj,color="black", linestyle="dotted", \\\n                   x_inline=True,y_inline=True)\n    gl.xlocator = mticker.FixedLocator(lons)\n    gl.ylocator = mticker.FixedLocator(lats)\n    iax.set_xticks(np.arange(-170,-60,4),crs=crs.PlateCarree())\n    iax.set_yticks(np.arange(10,70,4),crs=crs.PlateCarree())\n    lon_formatter = LongitudeFormatter()\n    lat_formatter = LatitudeFormatter()\n    iax.xaxis.set_major_formatter(lon_formatter)\n    iax.yaxis.set_major_formatter(lat_formatter)\n    \n    for tick in iax.xaxis.get_major_ticks():\n                tick.label.set_fontsize(10)\n    for tick in iax.yaxis.get_major_ticks():\n                tick.label.set_fontsize(10)\n\n    iax.set_extent([-125,-117.5,35,43])\n\nplot_hatch = ax1.pcolormesh(lon_wrf_d01,lat_wrf_d01,p1_d01[:-1,:-1],\\\n                hatch="/"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\nplot_hatch = ax1.pcolor(lon_wrf_d01,lat_wrf_d01,p3_d01[:-1,:-1],\\\n                hatch="x"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\n\nplot_hatch = ax2.pcolormesh(lon_wrf_d02,lat_wrf_d02,p1_d02[:-1,:-1],\\\n                hatch="/"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\nplot_hatch = ax2.pcolor(lon_wrf_d02,lat_wrf_d02,p3_d02[:-1,:-1],\\\n                hatch="x"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\n\nplot_hatch = ax3.pcolormesh(lon_wrf_d03,lat_wrf_d03,p1_d03[:-1,:-1],\\\n                hatch="/"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\nplot_hatch = ax3.pcolor(lon_wrf_d03,lat_wrf_d03,p3_d03[:-1,:-1],\\\n                hatch="x"*5,alpha=0,shading=\'flat\')\nmatplotlib.rcParams[\'hatch.linewidth\'] = 0.3\n\n#Get positions for cbs\nax1_pos = ax1.get_position()\nax2_pos = ax3.get_position()\n\ncbar_ax = fig.add_axes([ax1_pos.x0, ax1_pos.y0-0.047, ax2_pos.x1 - ax1_pos.x0, 0.015]) # bottom left corner x,y, and width, height\ncb = P.colorbar(cs,cax=cbar_ax,extend="both",orientation="horizontal")\ncb.ax.tick_params(labelsize=15)\n')


# In[ ]: