#!python3
# get KE geostrphic from the MEAN SSH
# based on Xiaobiao's code and my IDL code

import sys
prfx='/discover/nobackup/projects/gmao/cal_ocn/abozec1/PYTHON/'
sys.path.append(prfx)

import os
import proplot as plot
import netCDF4 as nc
import numpy as np

from myutilities.grid_mom6 import get_MOM6grid
from hycom.info import read_field_names,read_field_grid_names
from hycom.io import read_hycom_fields, read_hycom_grid,sub_var2
from myutilities.density import sigma2_hycom
from myutilities.tvplus import tvplus
from myutilities.plot_util import interp2plot,interp2plot2
from myutilities.get_cmap import  get_cmapeke, get_cmapsshvar,get_cmapjet40
from matplotlib.colors import BoundaryNorm
from scipy import ndimage


## get area grid
iog=prfx+'../MOM6-expt/GLBb0.08/RLX_SSH/YEARLY/'
file='regional.grid.a'
idm=4500 ; jdm=3298
kdm=41
i1=2574 ; i2=3203
j1=1898 ; j2=2335
idm1=i2-i1 ; jdm1=j2-j1
print ('idm1:',idm1, ' jdm1:',jdm1)
fieldg=['plon','plat','pscx','pscy','cori']
grid_field= read_hycom_grid(iog+file, fieldg)
plon=grid_field['plon'][j1:j2,i1:i2]
plat=grid_field['plat'][j1:j2,i1:i2]

pscx=grid_field['pscx'][j1:j2,i1:i2]
pscy=grid_field['pscy'][j1:j2,i1:i2]
cori=grid_field['cori'][j1:j2,i1:i2]
print('cori:',cori.shape)
## get grid and mask
plon[plon <= -180.]=plon[plon <= -180.]+360.
#plon[plon <= 0.]=plon[plon <= 0.]+360.
print('Model grid read! ')

# set path
io_data = prfx+'../MOM6-expt/GLBb0.08/MONTHLY_Z/'
year='2010_2014'
file1='ocnm_'+year+'_SSH2.nc'
ds1=nc.Dataset(io_data+file1)
sshm=ds1['SSH'][0,j1:j2,i1:i2]
ssh=sshm.data
ssh[ssh > 1e10] = np.nan 
print('SSH:',ssh.shape)

# cal geostrophic vel
g=9.806
ug = np.zeros([jdm1,idm1])
vg = ug
print('ug:',ug.shape)
print(ug[1:jdm1-1,:].shape, ssh[1:jdm1-1,:].shape, ssh[0:jdm1-2,:].shape,pscy[1:jdm1-1,:].shape, cori[1:jdm1-1,:].shape )

#;ug[*, 1:jpj-1] = -g*(ssh[*, 1:jpj-1] - ssh[*, 0:jpj-2])/(e2t[*, 1:jpj-1]*cori[*, 1:jpj-1])
# MOM6 grid versus HYCOM grid: MOM6 start at p-point, HYCOM starts at q-point
ug[0:jdm1-2,:] = -g*(ssh[1:jdm1-1,:] - ssh[0:jdm1-2,:])/(pscy[0:jdm1-2,:]*cori[0:jdm1-2,:])

#;vg[1:jpi-1, *] =  g*(ssh[1:jpi-1, *] - ssh[0:jpi-2, *])/(e1t[1:jpi-1, *]*cori[1:jpi-1, *])
vg[:,0:idm1-2] = -g*(ssh[:,1:idm1-1] - ssh[:,0:idm1-2])/(pscx[:,0:idm1-2]*cori[:,0:idm1-2])

#ug[*, 1:jpj-1] = -4*g*(ssh[*, 1:jpj-1] - ssh[*, 0:jpj-2])/((e2t[*,0:jpj-2]+e2t[*,1:jpj-1])*(cori[*,0:jpj-2]+cori[*,1:jpj-1]))
ug[0:jdm1-2,:] = -4.0*g*(ssh[1:jdm1-1,:] - ssh[0:jdm1-2,:])/((pscy[0:jdm1-2,:]+pscy[1:jdm1-1,:])*(cori[0:jdm1-2,:]+cori[1:jdm1-1,:]))


#vg[1:jpi-1, *] =  4*g*(ssh[1:jpi-1, *] - ssh[0:jpi-2, *])/((e1t[0:jpi-2,*]+e1t[1:jpi-1,*])*(cori[0:jpi-2,*]+cori[1:jpi-1,*]))
vg[:,0:idm1-2] = -4.0*g*(ssh[:,1:idm1-1] - ssh[:,0:idm1-2])/((pscx[:,0:idm1-2]+pscx[:,1:idm1-1])*(cori[:,0:idm1-2]+cori[:,1:idm1-1]))


# Needs smoothing ?
#ugs=smooth(ug,3,/nan)
#print, 'smoothing ug done! '
#vgs=smooth(vg,3,/nan)
#print, 'smoothing vg done! '


ke_geo=0.5*(ug**2+vg**2)
grid_x,grid_y=np.mgrid[-80.0:-30.:0.08,30.0:55.0:0.08]
ke_interp=interp2plot2(ke_geo,plon,plat,grid_x,grid_y,method='nearest')
print('ke interpolated ')

# colorbar eke
#cmap='coolwarm'
cmapeke=get_cmapeke()
N=cmapeke.colors.shape[0]
cmap=cmapeke

vmin=0.;vmax=3000.
levels_to_draw=np.linspace(vmin,vmax,N+1)

# for pcolormesh
norm0 = BoundaryNorm(levels_to_draw, ncolors=N, clip=True)
kem=np.nan_to_num(ke_interp*1e4,nan=0.)
y0 = ndimage.uniform_filter(kem,size=4)


# zoom North Atlantic
lonlim_na=(-100,0) ; latlim_na=(10,70)

# zoom GS
lonlim_gs=(280, 330) ; latlim_gs=(30, 55)

# choose your domain
lonlim=lonlim_gs ; latlim=latlim_gs

# plot the field
plot.rc.reso = 'med'  # use higher res for zoomed in geographic features
#proj = plot.Proj('mill', lonlim=(-99, -78), latlim=(18, 32), basemap=True)
#fig, axs = plot.subplots(nrows=1,ncols=2,proj='cyl',refwidth='4.5in')
fig, axs = plot.subplots(nrows=1,ncols=1,proj='mill',refheight='6cm',refwidth='10cm',tight=True)

# KE
axs[0].format(latlines=5, lonlines=5, labels=True,land=True,coast=True,landcolor='w',\
              title='KE geost. flow (cm2/s2) GLBb0.08 MOM6 '+year,lonlim=lonlim, latlim=latlim)
m = axs[0].pcolormesh(grid_x,grid_y, y0, cmap=cmap, norm=norm0, \
                               shading='auto', extend='neither')

#m = axs[0].contourf(plon, plat, ke*1e4, cmap=cmap, extend='neither',\
#                   levels=levels_to_draw)
axs[0].colorbar(m,loc='r',ticks=500.)

## EKEM
#axs[1].format(latlines=5, lonlines=5, labels=True,land=True,coast=True,landcolor='w',\
#              title='EKE (cm2/s2) GLBb0.08 MOM6 '+year,lonlim=lonlim, latlim=latlim)
#m = axs[1].pcolormesh(plon, plat, y, cmap=cmap, norm=norm, \
#                               shading='auto', extend='neither')
#
##m = axs[0].contourf(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], eke_model, cmap=cmap, extend='neither',\
##                    levels=levels_to_draw)
#axs[1].colorbar(m,loc='r',ticks=500.)

ps_dir=prfx+'/PNG/GLBb0.08/'
file_ps='glb0.08_GS_kegeo_'+year+'.png'
print(file_ps)
fig.savefig(ps_dir+file_ps,dpi=200,\
            facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps
#plot.close()


plot.show() ## interactive (block=False), otherwise ()
#plot.show(block=False)