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

# In[1]:


import xarray as xr
import gcsfs
import intake
import numpy as np
import matplotlib
#import cmocean
import stats
import stats as st
from matplotlib import pyplot as plt
get_ipython().magic(u'matplotlib inline')
import ecco_v4_tools as ecco
# Import plotting libraries
import importlib
import llcmapping
importlib.reload(llcmapping)
from llcmapping import LLCMapper
importlib.reload(ecco)
importlib.reload(ecco.tile_plot_proj)
import glob


# In[2]:


matplotlib.rcParams.update({'font.size': 16})
matplotlib.rcParams.update({'axes.titlesize': 16})
matplotlib.rcParams.update({'figure.figsize': (10,8)})
matplotlib.rcParams.update({'lines.linewidth': 2})
matplotlib.rcParams.update({'legend.fontsize': 18})
matplotlib.rcParams.update({'mathtext.fontset': 'cm'})
matplotlib.rcParams.update({'ytick.major.size': 3})
matplotlib.rcParams.update({'axes.labelsize': 16})
matplotlib.rcParams.update({'ytick.labelsize': 16})
matplotlib.rcParams.update({'xtick.labelsize': 16})


# In[3]:


fin = '/glade/work/patrizio/ECCO/' 
fout = '/glade/scratch/patrizio/figs/'


# In[4]:


fcoords = glob.glob(fin + 'fields/*coords.nc')[0]
fsnp = glob.glob(fin + 'fields/*snp.nc')
fvars = set(glob.glob(fin + 'fields/*.nc')) - set(glob.glob(fin + 'fields/*coords.nc')) - set(glob.glob(fin + 'fields/*snp.nc'))
fvars = list(fvars)
fvars 


# In[5]:


coords = xr.open_dataset(fcoords)
coords


# In[6]:


ntchunk=288

ds_snp = xr.open_mfdataset(fsnp,concat_dim=None)
ds_snp = ds_snp.chunk({'time':ntchunk-1})


# In[7]:


ds = xr.open_mfdataset(fvars,concat_dim=None)
ds


# In[8]:


ds_MXLDEPTH = xr.open_zarr(fin + 'ecco-data/MXLDEPTH', chunks={'time':ntchunk})
mxldepth = ds_MXLDEPTH.MXLDEPTH


# In[9]:


mxldepth


# For some reason the mixed layer depth coordinate indices are displaced by +1 in relation to the ECCO data stored on Pangeo. The coordinates need to be matched for future calculations. 

# In[10]:


mxldepth.coords['i'] = coords['i']
mxldepth.coords['j'] = coords['j']


# Calculate climatological mean mixed layer depth. We will be using this later to mask grid points outside of the mixed layer. 

# In[11]:


mxldepth_clim=mxldepth.mean(dim='time').load()
#mxldepth_clim=mxldepth.mean(dim='time').persist()


# Make a mask of points outside the ocean mixed layer:

# In[12]:


mxlpoints = np.abs(coords['Z']) <= mxldepth_clim


# In[ ]:





# In[13]:


# Flag for low-pass filtering
lowpass=True

# Filter requirements
order = 5
fs = 1     # sample rate, (cycles per month)
Tn = 12*3.
cutoff = 1/Tn  # desired cutoff frequency of the filter (cycles per month)

# Face numbers to analyze
# 0: Southern Ocean (Atlantic)
# 1: South Atlantic Ocean / Africa 
# 2: East North Atlantic / Europe
# 3: Southern Ocean (Indian)
# 4: Indian Ocean
# 5: Asia
# 6: Arctic
# 7: North Pacific (central)
# 8: West South Pacific
# 9: Southern Ocean (West Pacific)
# 10: North America / West North Atlantic
# 11: East South Pacific / South America
# 12: Southern Ocean(East Pacific)
#facen = [5,7]

#Note: longitude bounds can either be 0 < bounds < 360, or -180 < bounds < 180. 
#The only requirement is that the left longitude bound is less than the right bound 
#(along date line must use 0 < bounds < 360).
#(along prime meridian must use -180 < bounds < 180)

# Complete global 
#facen=[0,1,2,3,4,5,6,7,8,9,10,11,12]
#bnds = [0,359.9,-90,90]

#facen=[]
#bnds = [0,359.9,-90,90]

# Global (excluding polar regions)
#facen=[1,2,4,5,7,8,10,11]
#bnds = [0,359.9,-58,70]

#Southern Ocean (Atlantic)
#facen=[0]
#bnds = [-20,20,-58,-90]

#1: South Atlantic Ocean / Africa
#facen=[1]
#bnds = [-38,30,-58,10]

#2: East North Atlantic 
#facen=[2]
#bnds = [-38,30,10,70]

#3: Southern Ocean (Indian)
#facen=[3]
#bnds = [60,143,-58,-90]

#4: Indian Ocean
#facen=[4]
#bnds = [60,143,-58,10]

#7: North Pacific (central)
#facen=[7]
#bnds = [145,230,10,70]

#8: West South Pacific
#facen=[8]
#bnds = [145,230,-58,10]

#11: East South Pacific
#facen=[11]
#bnds = [-128,-38,-58,10]

#2, 10: North Atlantic
facen=[2,10]
bnds = [-80,0,10,70]

#5,7,10: North Pacific
#facen=[5,7,10]
#bnds = [100,270,10,70]

#4,5,7,8,10,11: Pacific
#facen=[4,5,7,8,10,11]
#bnds = [100,300,-70,70]

#5,7,8,10,11: Tropical Pacific
#facen=[5,7,8,10,11]
#bnds = [145,290,-15,15]

#5,7: KOE
#facen=[5,7]
#bnds = [120,180,15,60]


# In[14]:


rho0 = 1029 #sea-water density (kg/m^3)
c_p = 3994 #sea-water heat capacity (J/kg/K)


# In[15]:


coords=coords.isel(face=facen)


# In[16]:


# Vertical grid spacing
drF = coords.drF
hFacC = coords.hFacC
#rA = coords.rA.isel(face=facen).load()
#vol = drF*hFacC*rA.load()


# In[17]:


c_o = rho0*c_p*drF*hFacC


# In[18]:


T = ds_snp.T.isel(face=facen)
adv_ConvH = ds.adv_ConvH.isel(face=facen)
dif_ConvH = ds.dif_ConvH.isel(face=facen)
forcH = ds.forcH.isel(face=facen)


# In[19]:


dt = coords.time_snp[1:].load()
dt = dt.rename({'time_snp': 'time'})
# delta t in seconds. Note: divide by 10**9 to convert nanoseconds to seconds
dt.values = [float(t)/10**9 for t in np.diff(coords.time_snp)]

# time axis of dt should be the same as of the monthly averages
dt.time.values = coords.time[1:-1].values


# In[20]:


lons = coords.XC
lats = coords.YC


# In[21]:


T_anom, T_clim = st.anom(T) 
C_adv_anom, C_adv_clim = st.anom(adv_ConvH)
C_dif_anom, C_dif_clim = st.anom(dif_ConvH)
C_forc_anom, C_forc_clim = st.anom(forcH)
totalH_anom = C_adv_anom + C_dif_anom + C_forc_anom


# In[22]:


T_anom = T_anom.chunk({'time':ntchunk-1})
C_adv_anom = C_adv_anom.chunk({'time':ntchunk})
C_dif_anom = C_dif_anom.chunk({'time':ntchunk})
C_forc_anom = C_forc_anom.chunk({'time':ntchunk})


# In[ ]:


if lowpass:
    
    T_anom = T_anom.chunk({'time':288, 'j':10, 'i':10})
    
    C_adv_anom = C_adv_anom.chunk({'time':288, 'j':10, 'i':10})
    C_dif_anom = C_dif_anom.chunk({'time':288, 'j':10, 'i':10})
    C_forc_anom = C_forc_anom.chunk({'time':288, 'j':10, 'i':10})
    
    T_anom  = stats.butter_lowpass_filter_xr(T_anom, cutoff, fs, order)
    
    C_adv_anom  = stats.butter_lowpass_filter_xr(C_adv_anom, cutoff, fs, order)
    C_dif_anom = stats.butter_lowpass_filter_xr(C_dif_anom, cutoff, fs, order)
    C_forc_anom  = stats.butter_lowpass_filter_xr(C_forc_anom, cutoff, fs, order)
    
    totalH_anom = C_adv_anom + C_dif_anom + C_forc_anom


# In[ ]:


get_ipython().magic(u'time T_anom.load()')
get_ipython().magic(u'time C_adv_anom.load()')
get_ipython().magic(u'time C_dif_anom.load()')
get_ipython().magic(u'time C_forc_anom.load()')


# In[ ]:


tendH_perMonth = (T_anom.shift(time=-1)-T_anom)[:-1]


# In[ ]:


# Make sure time axis is the same as for the monthly variables
tendH_perMonth.time.values = coords.time[1:-1].values

# Convert tendency from 1/month to 1/s
tendH_perSec = tendH_perMonth/dt
tendH_perSec = tendH_perSec.transpose('face','time', 'k', 'j', 'i')


# In[ ]:


# Define tendH array with correct dimensions
tendH_anom = xr.DataArray(np.nan*np.zeros([len(facen),np.shape(tendH_perSec)[1]+2,50,90,90]),
                     coords={'face': facen, 'time': range(np.shape(tendH_perSec)[1]+2),'k': np.array(range(0,50)),
                             'j': np.array(range(0,90)),'i': np.array(range(0,90))},dims=['face', 'time','k', 'j','i'])

tendH_anom.time.values = coords.time.values


# In[ ]:


tendH_anom


# In[ ]:


tendH_anom.nbytes/1e9


# In[ ]:


# Add coordinates#
tendH_anom['XC'] = lons
tendH_anom['YC'] = lats
tendH_anom['Z'] = coords.Z

# Total tendency (degC/s)
tendH_anom.values[:,1:-1,:] = tendH_perSec.values
get_ipython().magic(u'time tendH_anom.load()')
#%time tendH.persist()


# In[ ]:


# Convert from degC/s to W/m^2
tendH_anom = c_o*tendH_anom
tendH_anom = tendH_anom.transpose('time','face', 'k', 'j', 'i')


# In[ ]:


face=0
k = 0
j = 15
i = 15

plt.figure(figsize=(14,10))
plt.subplot(2, 1, 1)
plt.plot(tendH_anom.time, tendH_anom.isel(face=face,k=k,j=j,i=i), lw=4, color='K', marker='.',label='total tendency')
plt.plot(C_forc_anom.time, C_forc_anom.isel(face=face,k=k,j=j,i=i), lw=2, color='C0', marker='.',label='forcing')
plt.plot(C_adv_anom.time, C_adv_anom.isel(face=face,k=k,j=j,i=i), lw=2, color='C1', marker='.',label='advection')
plt.axhline(0,color='k',lw=1)
plt.plot(C_dif_anom.time, C_dif_anom.isel(face=face,k=k,j=j,i=i), lw=2, color='C2',label='diffusion')
plt.setp(plt.gca(), 'xticklabels',[])
plt.legend(loc='best',frameon=False,fontsize=14)

plt.subplot(2, 1, 2)
plt.plot(totalH_anom.time, totalH_anom.isel(face=face,k=k,j=j,i=i), lw=4, color='red', marker='.',label='RHS')
plt.plot(tendH_anom.time, tendH_anom.isel(face=face,k=k,j=j,i=i), lw=2, color='blue', marker='.',label='LHS')
plt.plot(tendH_anom.time, (totalH_anom-tendH_anom).isel(face=face,k=k,j=j,i=i), lw=2, color='k', marker='.',label='RHS - LHS')
plt.legend(loc='best',frameon=False,fontsize=14)
plt.savefig(fout + 'sstbudget_anom_ts.png')


# In[ ]:


T_var = T_anom.var(dim='time')
get_ipython().magic(u'time T_var.load()')
#%time T_var.persist()


# In[ ]:


tendH_anom = tendH_anom/c_o


# In[ ]:


#tendH_anom = tendH_anom.transpose('time','face', 'k', 'j', 'i')
cov_adv = st.cov(tendH_anom, C_adv_anom)
cov_dif = st.cov(tendH_anom, C_dif_anom)
cov_forc = st.cov(tendH_anom, C_forc_anom)


# In[ ]:


cov_adv.nbytes/1e9


# In[ ]:


get_ipython().magic(u'time cov_adv.load()')
get_ipython().magic(u'time cov_dif.load()')
get_ipython().magic(u'time cov_forc.load()')


# In[ ]:


deltat = dt.mean()
deltat.compute()


# In[ ]:


r_1 = st.cor(T_anom, T_anom,lagx=1).compute()
r_1


# In[ ]:


fac = (deltat**2/(2*c_o*(1-r_1)))
fac.load()


# In[ ]:


T_var_sum = fac*(cov_adv + cov_dif + cov_forc)


# In[ ]:


get_ipython().magic(u'time T_var_sum.load()')
#%time T_var_sum.persist()


# In[ ]:


mapper = LLCMapper(coords)


# In[ ]:


k=0
mapper(T_var.isel(k=k), bnds=bnds, cmap='cubehelix_r', vmin=0,vmax=1.0)
mapper(T_var_sum.isel(k=k), bnds=bnds, cmap='cubehelix_r', vmin=0,vmax=1.0)


# The temperature variance budget is clearly balanced! Let's take a look at the contribution due to each term.

# In[ ]:


T_var_adv = fac*cov_adv
T_var_dif = fac*cov_dif
T_var_forc = fac*cov_forc


# In[ ]:


vmin=-1.0
vmax=1.0
sstmax=1.6
if lowpass:
   sstmax=0.5
   vmin=-0.5
   vmax=0.5


# ### Contributions to temperature variance from advection, diffusion and surface forcing

# In[ ]:


k=0
mapper(T_var_sum.isel(k=k), bnds=bnds, cmap='cubehelix_r', vmin=0,vmax=sstmax)
plt.title(r'temperature variance (K$^2$)')
plt.savefig(fout + 'Tvar_sum.png')
mapper(T_var_adv.isel(k=k), bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'advective contribution (K$^2$)')
plt.savefig(fout + 'Tvar_adv.png')
mapper(T_var_dif.isel(k=k), bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'diffusive contribution (K$^2$)')
plt.savefig(fout + 'Tvar_dif.png')
mapper(T_var_forc.isel(k=k), bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'surface forcing contribution (K$^2$)')
plt.savefig(fout + 'Tvar_forc.png')


# ### Contributions to ocean mixed layer temperature variance from advection, diffusion and surface forcing

# In[ ]:


mxlpoints = mxlpoints.isel(face=facen)
delz = drF*hFacC
delz=delz.where(mxlpoints)
delz_sum = delz.sum(dim='k')


# In[ ]:


mxlpoints


# In[ ]:


weights = delz/delz_sum


# In[ ]:


T_var_mxl = (weights*T_var).where(mxlpoints).sum(dim='k')


# In[ ]:


T_var_adv_mxl = (weights*T_var_adv).where(mxlpoints).sum(dim='k')
T_var_dif_mxl = (weights*T_var_dif).where(mxlpoints).sum(dim='k')
T_var_forc_mxl = (weights*T_var_forc).where(mxlpoints).sum(dim='k')


# In[ ]:


T_var_sum_mxl = T_var_adv_mxl + T_var_dif_mxl + T_var_forc_mxl


# In[ ]:


#f, axes = plt.subplots(2,2,figsize=(16,12))
#f.tight_layout()
mapper(T_var_sum_mxl, bnds=bnds, cmap='cubehelix_r', vmin=0,vmax=sstmax)
plt.title(r'temperature variance (K$^2$)')
plt.savefig(fout + 'Tmxlvar_sum.png')

mapper(T_var_adv_mxl, bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'advective contribution (K$^2$)')
plt.savefig(fout + 'Tmxlvar_adv.png')

mapper(T_var_dif_mxl, bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'diffusive contribution (K$^2$)')
plt.savefig(fout + 'Tmxlvar_dif.png')

mapper(T_var_forc_mxl, bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'surface forcing contribution (K$^2$)')
plt.savefig(fout + 'Tmxlvar_forc.png')


# In[ ]:


#mapper(T_var_sum_mxl, bnds=bnds, cmap='cubehelix_r', vmin=0,vmax=1.0)
#plt.title(r'temperature variance (K$^2$)')
#plt.savefig(fout + 'Tmxlvar_sum.png')
mapper(T_var_adv_mxl + T_var_dif_mxl, bnds=bnds, cmap='RdBu_r', vmin=vmin,vmax=vmax)
plt.title(r'ocean dynamics (advective + diffusive) contribution (K$^2$)')
plt.savefig(fout + 'Tmxlvar_ocndyn.png')
#mapper(T_var_forc_mxl, bnds=bnds, cmap='RdBu_r', vmin=-1.0,vmax=1.0)
#plt.title(r'surface forcing contribution (K$^2$)')
#plt.savefig(fout + 'Tmxlvar_forc.png')


# In[ ]:




