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

# ### Simple example
# This code shows as simple way to:
# - Read in a .nc file using xarray
# - Calculate annual values
# - Plot global time series
# - Plot spatial maps

# In[1]:


import xarray as xr
import cf_units as cf
import esmlab
from ctsm_py import utils
import numpy as np
import pandas as pd
from netCDF4 import num2date

# some resources for plotting
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
#import hvplot.xarray
#import holoviews as hv
#hv.extension('bokeh')
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


## User defined options
model = 'CLM5_GSWP3'
var  = 'GPP'
suff = 'lnd/proc/tseries/month_1/'
path = '/glade/p/cgd/tss/people/oleson/CLM_LAND_ONLY_RELEASE/CLM5/clm50_r270_1deg_GSWP3V1_iso_newpopd_hist/'+suff
file = path + 'clm50_r270_1deg_GSWP3V1_iso_newpopd_hist.clm2.h0.'+ var +'.185001-201412.nc'
file


# #### Read in your data

# In[3]:


# ?? Wrapping the time_set_mid utility corrects time being read in by xr.open_dataset, other options ??
# ?? Would be nice to build a way to read in multiple models here, maybe build on Keith L's Example ??
# data sets have lots of coords, variables, and attributes
#ds   = xr.open_dataset(file, decode_times=True)
ds   = utils.time_set_mid(xr.open_dataset(file, decode_times=True), 'time')
print(ds)
area = ds.area
landfrac = ds.landfrac
da = ds[var]    # strips out the extra coords and variables from the dataset
dataset = da.to_dataset()  # convert the previously created unnamed DataArray to Dataset with a variable named "array"


# #### Make weights for global sum

# In[4]:


# ?? Would be nice to use cfunits  here ??
landUp = area * landfrac * 1e6      # area in km2, not m2
orig_units = cf.Unit(landUp.attrs['units'])
target_units = cf.Unit('m^2')
landUp.attrs['units'] = target_units# rename units
area_wgt = landUp / landUp.sum()    # weighting for each grid cell
spy = 365 * 24 * 3600               # Convert  to annual fluxes (gC/m2/y), from gc/m2/s 


# #### Convert monthly to annual

# In[5]:


# apply iterates through all variables in a dataset and applies the function to each variable.
# ?? Would be nice to do this differently for fluxes & stocks ??
mean_dataset = dataset.apply(utils.weighted_annual_mean) * spy
mean_dataset.attrs['units'] = 'gC m-2 y-1'
print(mean_dataset)


# #### Calculate global sums

# In[6]:


# ?? This could also be done regionally (see Paul's comments on google doc) ??
mean_wgt = mean_dataset * landUp
ann_glob = mean_wgt.sum(dim=('lat','lon')) * 1e-15 #convert to Pg/y
ann_glob.attrs['units'] = 'Pg C y^-1'
print(ann_glob)


# -----------------------------
# ### Make a line plot

# In[7]:


#quick plot , the `;` on the last line supresses text after executing the cell
plt.plot(ann_glob.coords['time.year'], ann_glob[var], '-' );


# In[8]:


# Create a nicer figure
fig = plt.figure(figsize=(10, 6))

# Ask, out of a 1x1 grid, the first axes.
ax = fig.add_subplot(1, 1, 1)

# Plot times as x-variable and temperatures as y-variable
ax.plot(ann_glob.coords['time.year'], ann_glob[var], label=model)
# Add some labels to the plot
ax.set_xlabel('Time')
ax.set_ylabel(var+' (' + ann_glob.attrs['units'] +')')
ax.set_title(var, fontdict={'size':16})
ax.legend(loc='upper left');


# ------------------------
# ### Now make a map!

# In[9]:


# making a data array, because cyclic_dataarray currently set up that way
x1= mean_dataset.isel(time=slice(0,10)).mean('time')[var] # Climatology for first 10 years of data
x = mean_dataset.isel(time=slice(-10,None)).mean('time')[var] # Climatology for last 10 years of data
xdiff = x - x1

xinit     = x1 * landUp
xfin      = x  * landUp
xdelta    = xdiff * landUp

glob_init   = xinit.sum(dim=('lat','lon')) * 1e-15 #convert to Pg/y
glob_fin    = xfin.sum(dim=('lat','lon')) * 1e-15 #convert to Pg/y
glob_change = xdelta.sum(dim=('lat','lon')) * 1e-15 #convert to Pg/y
print(glob_init, glob_fin, glob_change)

# Then add cyclic point, ?? is there a better way to do this ??
x1= utils.cyclic_dataarray(x1)  
x = utils.cyclic_dataarray(x)
xdiff = utils.cyclic_dataarray(xdiff)
lat = x.coords['lat']
lon = x.coords['lon']


# In[10]:


import matplotlib.colors as colors
import cartopy.feature as cfeature

norm = colors.Normalize(vmin=0., vmax=3.5e3)
norm(x.mean())

# modified from https://unidata.github.io/MetPy/latest/examples/Four_Panel_Map.html
crs = ccrs.Robinson()#central_longitude=305.0)

# Function used to create the map subplots
def plot_background(ax):
    ax.set_global()
    ax.coastlines()
    #ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5)
    return ax

levels = np.linspace(0, 3500, 11)
fig = plt.figure(figsize=(5, 3.5),constrained_layout=True);
fig, axarr = plt.subplots(nrows=2, ncols=1, figsize=(10, 6), constrained_layout=True,
                          subplot_kw={'projection': crs})

axlist = axarr.flatten()
for ax in axlist:
    plot_background(ax)
    ax.set_extent([-155, 60, -25, 60])  # I really don't understand how this works??
#    ax.set_extent([-150, 60, -25, 60])


# upper plot
cf0 = axlist[0].contourf(lon, lat, x1, levels=levels, norm=norm, cmap='gist_earth_r',
                 transform=ccrs.PlateCarree());
axlist[0].set_title('Initial '+model+' '+var, fontsize=18)
axlist[0].set_axis_off()

# lower plot
cf1 = axlist[1].contourf(lon, lat, x, levels=levels, norm=norm, cmap='gist_earth_r',
                 transform=ccrs.PlateCarree());
axlist[1].set_title('Final '+model+' '+var, fontsize=18);
axlist[1].set_axis_off()

# add common colorbar
cbar = fig.colorbar(cf1, ax=axlist.ravel().tolist(), 
                    orientation='vertical', shrink=0.80, pad=0)
#cbar.set_label('gC/m2/y', size=18, orientation='horizontal')
cbar.ax.tick_params(labelsize=16) 
ax = cbar.ax
ax.text(0.2,-0.1,mean_dataset.attrs['units'], size=16,rotation=0)
plt.show();


# In[11]:


# Now make a difference map
norm = colors.Normalize(vmin=-1000., vmax=1000.)
norm(xdiff.mean())
levels = np.linspace(-1e3, 1e3, 11)

ax = plt.axes(projection=ccrs.Robinson())
ax.set_global()
ax.coastlines()
ax.set_extent([-155, 60, -25, 60])  # I really don't understand how this works??

cs = ax.contourf(lon, lat, xdiff, levels=levels, norm=norm, cmap='BrBG',
                 transform=ccrs.PlateCarree());

# add colorbar.
cbar = plt.colorbar(cs, orientation='horizontal')
cbar.set_label('gC/m2/y')

ax.set_title(var + " change")
plt.show();


# In[ ]: