#!python3
## plot GLBb0.08 EKE from the Gulf Stream field from MOM6
# to avoid attribute error due to version in proplot ...
from importlib.metadata import version
from packaging.version import parse

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 xarray as xr
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,get_cmapjet40
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/'
file='regional.grid.a'
idm=4500 ; jdm=3298
kdm=41
i1=2574 ; i2=3203
j1=1898 ; j2=2335

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'][:,:]
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
area_gs=area[j1:j2,i1:i2]
## 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'
expt='10.8'
io_data = prfx+'../MOM6-expt/GLBb0.08/expt_'+expt+'/NC_LAYERS/'
file1='ocnm_'+year+'_SSH2.nc'
ds1=xr.open_dataset(io_data+file1)
# get 3 year  mean SSH from dailies 
ssh=ds1.SSH[0,j1:j2,i1:i2] # in m
# get 3 year mean of SSH^2 from daily squares
ssh2=ds1.SSH2[0,j1:j2,i1:i2] # in m^2
#ds1=nc.Dataset(io_data+file1)
#ssh=ds1['SSH'][0,j1:j2,i1:i2]
#ssh2=ds1['SSH2'][0,j1:j2,i1:i2]
print(ssh.shape)
print(ssh2.shape)

sshm=ssh.data
# remove domain mean
ssh_mean=np.nansum(sshm*area_gs)/np.nansum(area_gs)
print(ssh_mean)
ssh_ref=sshm.copy()
ssh_ref[:,:]=(ssh_ref[:,:]-ssh_mean)*100.
ssh2m=ssh2.data
sshm[np.isnan(sshm) == True]=0.
ssh2m[np.isnan(ssh2m) == True]=0.
#sshm=sshm
#ssh2m=ssh2m
#sshm_model=np.roll(sshm, ishift,axis=1)

# variability of SSH in cm
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)

# interpolate to regular grid 
grid_x,grid_y=np.mgrid[-80.0:-30.:0.08,30.0:55.0:0.08]
ssh_interp=interp2plot2(ssh_ref,plon,plat,grid_x,grid_y,method='nearest')
print('ssh interpolated ')
sshvar_interp=interp2plot2(ssh_std,plon,plat,grid_x,grid_y,method='nearest')
print('sshvar interpolated ')



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

# for pcolormesh
norm0 = BoundaryNorm(levels_ssh, ncolors=40, clip=True)
ssh_interp=np.nan_to_num(ssh_interp,nan=0.)
y0 = ndimage.uniform_filter(ssh_interp,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(sshvar_interp,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='mill',refheight='6cm',refwidth='10cm',tight=True)
#SSH mean
axs[0].format(latlines=5, lonlines=5, coast=True,labels=True,land=True,landcolor='w',landzorder=2,\
              title='SSH (cm) GLBb0.08 MOM6 '+expt+' '+year,lonlim=(-80,-30), latlim=(30, 54))
axs[0].contour(grid_x,grid_y,ssh_interp,\
              levels=np.linspace(-90,90,17),color='gray7',lw=0.5)
m = axs[0].pcolormesh(grid_x,grid_y,y0, cmap=cmap_ssh, norm=norm0, \
                               shading='auto', extend='neither')

axs[0].colorbar(m,loc='r',ticks=20.)

axs[1].format(latlines=5, lonlines=5, coast=True,labels=True,land=True,landcolor='w',\
              title='SSH variability (cm) GLBb0.08 MOM6 '+expt+' '+year,lonlim=(280, 330), latlim=(30, 54))
m = axs[1].pcolormesh(grid_x,grid_y, y, cmap=cmap_std, norm=norm, \
                               shading='auto', extend='neither')
axs[1].colorbar(m,loc='r',ticks=5.)

ps_dir=prfx+'/PNG/GLBb0.08/'
file_ps='glb0.08_GS_ssh-sshvar_'+year+'b_'+expt+'.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)