#!python3
## plot GLBb0.08 EKE from the Gulf Stream field from MOM6
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
from matplotlib.colors import BoundaryNorm
from scipy import ndimage


## change land color
plot.rc.landcolor=[191/256,159/256,127/256]


## get area grid
iog=prfx+'../MOM6-expt/GLBb0.08/RLX_SSH/YEARLY/'
file='regional.grid.a'
idm=4500 ; jdm=3298
i1=2388 ; i2=3611
j1=1473 ; j2=2573
idm1=i2-i1+1
jdm1=j2-j1+1

kdm=41
fieldg=['plon','plat','pscx','pscy']
grid_field= read_hycom_grid(iog+file, fieldg)
plon=grid_field['plon'][j1:j2,i1:i2]
plat=grid_field['plat'][j1:j2,i1:i2]
dx=grid_field['pscx'][j1:j2,i1:i2]
dy=grid_field['pscy'][j1:j2,i1:i2]
## get mask
#bathy=sub_var2(iog+'regional.depth.a',idm,jdm,1)
#mask=np.zeros([jdm,idm])+1.
#mask[np.isnan(bathy) == True ] = 0.
#plon[mask == 0] = 2**100
#plat[mask == 0] = 2**100
#
#area=dx*dy*mask
## get grid and mask
plon[plon <= -180.]=plon[plon <= -180.]+360.
#plon[plon <= 0.]=plon[plon <= 0.]+360.
print(plon[0,483])

#ishift= -483 ## plon=-179.89
#ishift=-1197 ## plon=360.0
#lon_model=np.roll(plon, ishift,axis=1)
#lat_model=np.roll(plat, ishift,axis=1)
#print(lon_model.shape,lat_model.shape)
#print('Model grid read! ')

# set path
year='2010_2014'
io_data = prfx+'../MOM6-expt/GLBb0.08/MONTHLY_Z/'
file1='ocnm_'+year+'_surf.nc'
fields=['KE','u','v']
layers=[0]
ds1=nc.Dataset(io_data+file1)
tmp_um=ds1['u'][0,0,j1:j2,i1:i2]
tmp_vm=ds1['v'][0,0,j1:j2,i1:i2]
print(tmp_um.shape)
print(tmp_vm.shape)

um=tmp_um.data
vm=tmp_vm.data
um1=np.roll(um,-1,axis=1) # i+1
vm1=np.roll(vm,-1,axis=0) # j+1
um[um > 1e10]=0. ; um1[um1 > 1e10]=0.
vm[vm > 1e10]=0. ; vm1[vm1 > 1e10]=0.

ke=ds1['KE'][0,0,j1:j2,i1:i2]
ke=ke.data
ke[ke > 1e10] = 0.0
print(ke.shape)
print(np.nanmin(ke), np.nanmax(ke))

# KE (cm2/s2) of the mean flow
ke_model=(0.125*((um+um1)**2 + (vm+vm1)**2))*1e4 # with u and v on the same T-grid


#ke_model=np.roll(ke*1e4, ishift,axis=1)
# EKE = Total KE - KE of the mean flow (cm2/s2)
eke=(ke - 0.125*((um+um1)**2 + (vm+vm1)**2))*1e4 # u and v on the T-grid
print(eke.shape)
print(np.nanmin(eke), np.nanmax(eke))
#eke_model=np.roll(eke, ishift,axis=1)

# 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_model,nan=0.)
y0 = ndimage.uniform_filter(kem,size=4)

# for pcolormesh
norm = BoundaryNorm(levels_to_draw, ncolors=N, clip=True)
ekem=np.nan_to_num(eke,nan=0.)
y = ndimage.uniform_filter(ekem,size=4)

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

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

# 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=2,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 of the mean flow (cm2/s2) GLBb0.08 MOM6 '+year,lonlim=lonlim, latlim=latlim)
m = axs[0].pcolormesh(plon, plat, 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_ke-eke_'+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)