#!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_cmapsshvar
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
kdm=41
fieldg=['plon','plat','pscx','pscy']
grid_field= read_hycom_grid(iog+file, fieldg)
plon=grid_field['plon'][:,:]
plat=grid_field['plat'][:,:]
dx=grid_field['pscx'][:,:]
dy=grid_field['pscy'][:,:]
## 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='2015_2017'
io_data = prfx+'../MOM6-expt/GLBb0.08/MONTHLY_Z/'
file1='ocnm_'+year+'_SSH2.nc'
ds1=nc.Dataset(io_data+file1)
ssh=ds1['SSH'][0,:,:]
ssh2=ds1['SSH2'][0,:,:]
print(ssh.shape)
print(ssh2.shape)

sshm=ssh.data
ssh2m=ssh2.data
sshm[sshm > 1e10]=0.
ssh2m[ssh2m > 1e10]=0.
sshm=sshm
ssh2m=ssh2m
sshm_model=np.roll(sshm, ishift,axis=1)

ssh_std=np.sqrt(ssh2m - sshm**2)*100.
print(ssh_std.shape)
print(np.nanmin(ssh_std), np.nanmax(ssh_std))
sshstd_model=np.roll(ssh_std, ishift,axis=1)

# colorbar eke
#cmap='coolwarm'
cmap_ssh='jet'
vmin=-100. ; vmax=100.
levels_ssh = np.linspace(vmin,vmax,41)

# for pcolormesh
norm0 = BoundaryNorm(levels_ssh, ncolors=250, clip=True)
sshm=np.nan_to_num(sshm_model*100,nan=0.)
y0 = ndimage.uniform_filter(sshm,size=4) # convert ot cm


cmap_std=get_cmapsshvar()
vmin=0.;vmax=40.
levels_to_draw=np.linspace(vmin,vmax,17)

# for pcolormesh
norm = BoundaryNorm(levels_to_draw, ncolors=16, clip=True)
sshstdm=np.nan_to_num(sshstd_model,nan=0.)
y = ndimage.uniform_filter(sshstdm,size=4)



# 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')
#SSH mean
axs[0].format(latlines=10, lonlines=10, labels=True,land=True,\
              title='SSH (cm) GLBb0.08 MOM6 '+year,lonlim=(280, 330), latlim=(30, 55))
axs[0].contour(lon_model[0:jdm-1,:],lat_model[0:jdm-1,:],sshm_model*100,\
              levels=np.linspace(-90,90,17),color='gray7',lw=0.5)
m = axs[0].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y0, cmap=cmap_ssh, norm=norm0, \
                               shading='auto', extend='neither')
axs[0].colorbar(m,loc='b',ticks=20.)

axs[1].format(latlines=10, lonlines=10, labels=True,land=True,\
              title='SSH variability (cm) GLBb0.08 MOM6 '+year,lonlim=(280, 330), latlim=(30, 55))
m = axs[1].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y, cmap=cmap_std, norm=norm, \
                               shading='auto', extend='neither')
axs[1].colorbar(m,loc='b',ticks=5.)

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


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