%reset
%matplotlib inline

import netCDF4 as nc
import cmocean
import numpy as np
import xarray as xr
import xrft
from xgcm import Grid
import matplotlib.pyplot as plt
import scipy.interpolate
from dask.diagnostics import ProgressBar
from matplotlib import colors, ticker
import pandas as pd

%cd /glade/p/univ/unyu0004/eyankovsky/NeverWorld_analysis/
forcing=xr.open_dataset('NW2_forcing.nc', decode_times=False) #This is for the 1/8 degree case

#1/32 degree resolution, surface diagnostics are for uppermost 5 meters.
%cd /glade/p/univ/unyu0004/gmarques/NeverWorld2/baselines/nw2_0.03125deg_N15_baseline_hmix5
fs = xr.open_dataset('static.nc', decode_times=False)
os = xr.open_dataset('sub01/ocean.stats.nc', decode_times=False)
av = xr.open_dataset('averages_00002702.nc', decode_times=False)  #20 5-day averages
%cd /glade/p/univ/unyu0004/eyankovsky/NeverWorld_analysis/

##Increasing BS
# %cd /glade/p/univ/unyu0004/eyankovsky/MEKE_testing/MEKE_test16/
# av = xr.open_dataset('averages_00031502.nc', decode_times=False)  
# fs = xr.open_dataset('static.nc', decode_times=False)

%cd /glade/scratch/eyankovsky/Backscatter_tests_with_Scott/Half_degree/Jansen_EBT_test2/
av= xr.open_dataset('averages_00031502.nc', decode_times=False)  
fs = xr.open_dataset('static.nc', decode_times=False)

# %cd /glade/scratch/eyankovsky/Backscatter_tests_with_Scott/Half_degree/Jansen_EBT_test4/
# av= xr.open_dataset('averages_00031502.nc', decode_times=False)  
# fs = xr.open_dataset('static.nc', decode_times=False)

# %cd /glade/scratch/eyankovsky/Backscatter_tests_with_Scott/Half_degree/Jansen_EBT_test5
# av= xr.open_dataset('averages_00031502.nc', decode_times=False)  
# fs = xr.open_dataset('static.nc', decode_times=False)

# %cd /glade/scratch/eyankovsky/Backscatter_tests_with_Scott/Half_degree/Bachman_EBT_test4
# av= xr.open_dataset('averages_00031502.nc', decode_times=False)  
# fs = xr.open_dataset('static.nc', decode_times=False)

dx = (fs.dxCu.mean('xq')**2+fs.dyCu.mean('xq')**2)**0.5

lat= (av.yh)     
lon= (av.xh)     

u= av.u[:,0,:,:]-(av.u[:,0,:,:].mean('time')) #We're removing the temporally-averaged mean flow.
v= av.v[:,0,:,:]-(av.v[:,0,:,:].mean('time'))
grid= Grid(av, coords={'X': {'center': 'xh', 'outer': 'xq'}, #This tells the grid where xh, xq, yh, and yq are defined.
                        'Y': {'center': 'yh', 'outer': 'yq'}});
#Interpolating u,v to be on the same grid points:
dx= grid.interp(fs.dxCu,axis='X')
dy= grid.interp(fs.dyCv,axis='Y')
u = grid.interp(u,axis='X')
v = grid.interp(v,axis='Y')
print(u.shape,v.shape)
Ld = av.Rd1[:].mean('time')

#Plotting KE for whole domain
fig = plt.figure(figsize=(5, 7), dpi=200, facecolor='w', edgecolor='k')
ax1 = fig.add_axes([0.15,0.2,0.6,0.7]) 
ax2 = fig.add_axes([.78, 0.2, 0.03, 0.7])
ax1.set_facecolor((0.92, 0.92, 0.92))
ax1.set_xlabel('Longitude',fontsize=12)
ax1.set_ylabel('Latitude',fontsize=12)
ax1.tick_params(labelsize=12); 
ax2.tick_params(labelsize=12)


ax1.set_title('Surface EKE, 1/32$^\circ$',fontsize=12)

KE=(u**2+v**2)*0.5
plotted =ax1.pcolor(lon[:],lat[:],np.log10(KE[-1,:,:]),cmap=cmocean.cm.matter_r,vmin=-5,vmax=0)
ax1.set_aspect('equal', 'box')
cbar = plt.colorbar(mappable=plotted, cax=ax2, orientation = 'vertical',extend='both');
cbar.ax.tick_params(labelsize=12);
ax2.set_title('$log_{10}[m^2/s^2]$')
#
plt.show()

lat_full = fs['geolat']
lim = 5 #80 or 5
SSH = (av['h'][:,:,:,:]).sum('zl')-(av['h'][:,:,:,:]).sum('zl').mean('time');
print('done')
dSSHdx = SSH.differentiate('xh')*360/(6370000*2*np.pi*np.cos(lat_full*np.pi/180))
dSSHdy = SSH.differentiate('yh')*360/(6370000*2*np.pi)

numerator = (SSH**2).mean('time'); numerator = numerator.where(numerator>0)
denominator = ((abs(dSSHdy))**2+(abs(dSSHdx))**2).mean('time'); denominator = denominator.where(denominator>0e-20)

scale2 = ((numerator/denominator)**0.5).mean('xh')
scale2 = scale2[lim:-lim]/1000
scale2[0:10]=np.nan

dx

fig, ax = plt.subplots(figsize=(12,6),dpi=100)


plt.plot(scale_HR.yh,Ld[80:-80].mean('xh')/1000,'-',c='k',linewidth=3,label='$R_D$')
#plt.plot(ES_wavenumber.yh,1./ES_wavenumber,'-',c='k',linewidth=3,label='$R_E$ - spectral peak')
plt.plot(scale_HR.yh,scale_HR,'-',c='orchid',linewidth=2,label='$R_E$ for $1/32^\circ$')

plt.plot(scale1.yh,scale1,'-',color='mediumblue',linewidth=2,label='$R_E$ for $1/2^\circ$, Jansen et al. (2020)')

plt.plot(scale2.yh,scale2,'-',color='seagreen',linewidth=2,label='$R_E$ for $1/2^\circ$, our scheme')
#plt.plot(scale3.yh,scale3,'-',color='mediumturquoise',linewidth=2,label='$1/2^\circ$, Malte EBT Power 5')
#plt.plot(scale4.yh,scale4,'-',color='gold',linewidth=2,label='$1/2^\circ$, Scott EBT Power 4')
#plt.plot(dx.yh,dx,'-','c',[.5,.5,.5],linewidth=2,label='$R_E$ for $1/2^\circ$, our scheme')

plt.plot(dx.yh,dx/1000,'--',c=[.5,.5,.5],linewidth=2,label='Grid spacing, $1/2^\circ$')
plt.title('Energy Containing Scale Comparison',fontsize=20)
plt.ylim([0, 400]); plt.xticks(fontsize=16);
plt.xlim([-70, 70]); plt.yticks(fontsize=16);
plt.grid()
plt.xlabel('Latitude',fontsize=16); plt.ylabel('$[km]$',fontsize=14);
plt.legend(fontsize=14)

#ax.set_xlabel(r'k [1/km]')
#ax.set_ylabel(r'[m$^3$ s$^{-2}$]')

#COMPUTE velocity components for decomposing KE:
Nlayers=15

h  = av['h'].mean('time').fillna(0.)
Depth=h.sum('zl')
print('done part 1')
#Outputs are comprised of 100 5-day averages.
u  = av['u'];  u = grid.interp(u,axis='X')#.fillna(0.); 
ubar=u.mean('time')
print('done part 2')
v  = av['v'];  v = grid.interp(v,axis='Y')#.fillna(0.);
vbar=v.mean('time')

uprime=u-ubar
vprime=v-vbar

print(uprime.shape)
area_t=fs['area_t']

Depth.shape

# The mean BT and BC KE values:
ubar_TW=ubar*h/Depth; 
ubar_BT=ubar_TW.sum('zl')
ubar_BC=ubar-ubar_BT;

vbar_TW=vbar*h/Depth
vbar_BT=vbar_TW.sum('zl')
vbar_BC=vbar-vbar_BT;

KE_mean_BT=0.5*(vbar_BT**2+ubar_BT**2) #dim y,x
KE_mean_BC=(0.5*(vbar_BC**2+ubar_BC**2)*h/Depth).sum('zl') #dim y,x
print('done 1')

#Eddy BT and BC KE parts:
uprime_BT=(uprime*h/Depth).sum('zl')
uprime_BC=uprime-uprime_BT
vprime_BT=(vprime*h/Depth).sum('zl')
vprime_BC=vprime-vprime_BT
    
KE_eddy_BT=(0.5*(vprime_BT**2+uprime_BT**2)).mean('time')
KE_eddy_BC=((0.5*(vprime_BC**2+uprime_BC**2)*h/Depth).sum('zl')).mean('time')

print('done 2')
# #Total kinetic energy:
KE_BT_tot=KE_eddy_BT+KE_mean_BT
KE_BC_tot=KE_eddy_BC+KE_mean_BC
KE_tot=KE_BT_tot+KE_BC_tot


fraction_BT=KE_BT_tot/KE_tot
fraction_BT_mean=KE_mean_BT/KE_tot
fraction_BT_eddy=KE_eddy_BT/KE_tot

fraction_BC=KE_BC_tot/KE_tot#when you divide zero/zero it's a nan
fraction_BC_mean=KE_mean_BC/KE_tot
fraction_BC_eddy=KE_eddy_BC/KE_tot

KE_eddy = KE_eddy_BT+ KE_eddy_BC
KE_mean = KE_mean_BT+ KE_mean_BC

window=100

plt.figure(figsize=(10, 4), dpi= 100, facecolor='w', edgecolor='k')
#plt.plot(yh,np.nanmean(fraction_BT,axis=1),'-',c='#82cafc',linewidth=3)
plt.plot(fraction_BT.yh,np.convolve(np.nanmean(KE_eddy_BT/KE_eddy,axis=1), np.ones((window,))/window, mode='same'),'-',c='#ffa500',linewidth=3)

plt.grid()
plt.plot(fraction_BT.yh,np.convolve(np.nanmean(KE_eddy_BC/KE_eddy,axis=1), np.ones((window,))/window, mode='same'),'-',c=[0,0,.5],linewidth=3)


plt.legend(['BT eddy','BC eddy'],loc='upper right',fontsize=12,handlelength=4)
plt.plot([0, 0],[0, 1],'-',color=[0,0,0])
plt.xlim(-70, 70)
#plt.ylim(0,1);
plt.xlabel('Latitude')
plt.ylabel('Fraction')
plt.title('1/32$^\circ$ BT,BC EKE fractions')

window=20

plt.figure(figsize=(10, 4), dpi= 100, facecolor='w', edgecolor='k')
#plt.plot(yh,np.nanmean(fraction_BT,axis=1),'-',c='#82cafc',linewidth=3)
yh_center=av.yh[lim:-lim]
Eddy_BC=(KE_eddy_BC.mean('xh'))[lim:-lim];
Eddy_BT=(KE_eddy_BT.mean('xh'))[lim:-lim];

BT_BC=np.convolve((Eddy_BT/Eddy_BC)**0.5, np.ones((window,))/window, mode='same')
#LES_LD=np.convolve((LD_wavenumber/ES_wavenumber), np.ones((window,))/window, mode='same')
#LEC_LD=np.convolve((LD_wavenumber/centroid), np.ones((window,))/window, mode='same')
LEscale_LD=np.convolve((scale4/(Ld[lim:-lim].mean('xh')/1000)), np.ones((window,))/window, mode='same')

bound1=lim; bound2=-lim
plt.plot(yh_center[bound1:bound2],BT_BC[bound1:bound2],'-',c='k',linewidth=3)

# plt.plot(yh_center[bound1:bound2],LES_LD[bound1:bound2],'-',c='#FF5733',linewidth=3)
# plt.plot(yh_center[bound1:bound2],LEC_LD[bound1:bound2],'-',c='#5F9EA0',linewidth=3)
plt.plot(yh_center[bound1:bound2],LEscale_LD[bound1:bound2],'-',c='#FF5733',linewidth=3)


plt.legend(['(BT/BC EKE)$^{1/2}$','$R_E/R_D$'],loc='upper right',fontsize=18,handlelength=4)
#plt.plot([0, 0],[0, 1],'-',color=[0,0,0])
plt.ylim([0, 3]); plt.xticks(fontsize=18);
plt.xlim([-60, 60]); plt.yticks(fontsize=18);
plt.xlabel('Latitude',fontsize=18)
#plt.ylabel('Fraction')
#plt.title('1/32$^\circ$')
plt.grid()

plt.figure(figsize=(10, 4), dpi= 200, facecolor='w', edgecolor='k')
#plt.plot(yh,np.nanmean(fraction_BT,axis=1),'-',c='#82cafc',linewidth=3)
yh_center=av.yh[lim:-lim]
Eddy_BC=(KE_eddy_BC.mean('xh'))[lim:-lim];
Eddy_BT=(KE_eddy_BT.mean('xh'))[lim:-lim];
BT_BC=(Eddy_BT/Eddy_BC)**0.5
LE_LD=(LD_wavenumber*scale)  #Using TY06 scaling

bound1=80; bound2=-80
plt.scatter(LE_LD[bound1:bound2],BT_BC[bound1:bound2],c=yh_center[bound1:bound2], cmap='RdYlBu',marker='o',s=10)


#plt.xlim(0,5)
#plt.ylim(0,5);
plt.ylim([1e-1, 1e1])
plt.xlim([1e-1, 1e1])
plt.xscale('log')
plt.yscale('log')

plt.ylabel('(BT/BC EKE)$^{1/2}$',Fontsize=14)
plt.xlabel('$R_E/R_D$ Ratio',Fontsize=14)
plt.grid()
clb=plt.colorbar()
clb.ax.set_title('Lat.',fontsize=14)


coefficients = np.polyfit(np.log10(LE_LD[bound1:bound2]), np.log10(BT_BC[bound1:bound2]), 1)
polynomial = np.poly1d(coefficients)
log10_y_fit = polynomial(np.log10(LE_LD[bound1:bound2])) 
#
plt.plot([.3,3], [.3,3], 'k-',label='slope=1')  
plt.legend(fontsize=12)
ax = plt.axes()
ax.set_facecolor([0.75,0.75,0.75])
plt.gca().set_aspect('equal')