#!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_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/'
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
expt='10.8'
year='2015-2017'
io_data = prfx+'../MOM6-expt/GLBb0.08/expt_'+expt+'/MONTHLY_Z/'
file1='ocean_month_z_'+year+'.nc'
ds1=xr.open_dataset(io_data+file1)
tmp_um=ds1.u[0,0,:,0:idm]
tmp_vm=ds1.v[0,0,0:jdm-1,:]
#ds1=nc.Dataset(io_data+file1)
#tmp_um=ds1['u'][0,0,:,0:idm]
#tmp_vm=ds1['v'][0,0,0:jdm-1,:]
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[np.isnan(um) == True]=0. ; um1[np.isnan(um1) == True]=0.
vm[np.isnan(vm) == True]=0. ; vm1[np.isnan(vm1) == True]=0.

file1='ocean_month_z_'+year+'.nc'
ds1=xr.open_dataset(io_data+file1)
kemz=ds1.KE[0,0,:,0:idm]
kemz=kemz.data ; kemz[np.isnan(kemz) == True] = 0.0

# Ke of the mean flow (cm2/s2)
ke=0.125*((um+um1)**2 + (vm+vm1)**2)*1e4 ## On T-grid
ke_model=np.roll(ke, ishift,axis=1)

# EKE = total KE - KE of the mean flow (cm2/s2)
eke=(kemz-0.125*((um+um1)**2+(vm+vm1)**2))*1e4 # eke 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_model,nan=0.)
y = ndimage.uniform_filter(ekem,size=4)


domain='NA'
if (domain == 'NA'): 
  # zoom North Atlantic
  lonlim_na=(-100,0) ; latlim_na=(10,70)
  lonlim=lonlim_na ; latlim=latlim_na
  
  plot.rc.reso = 'med'  # use higher res for zoomed in geographic features
  fig, axs = plot.subplots(nrows=1,ncols=2,proj='mill',refwidth='10cm')
  
  # KE
  axs[0].format(latlines=10, lonlines=10, labels=True,land=True,coast=False,\
                title='KE of the mean flow (cm2/s2) GLBb0.08 MOM6 '+expt+' '+year,lonlim=lonlim, latlim=latlim)
  
  m = axs[0].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y0, cmap=cmap, norm=norm0, \
                                 shading='auto', extend='neither')
  
  axs[0].colorbar(m,loc='b',ticks=500.)
  
  
  # EKEM
  axs[1].format(latlines=10, lonlines=10, labels=True,land=True,coast=False,\
                title='EKE (cm2/s2) GLBb0.08 MOM6 '+expt+' '+year,lonlim=lonlim, latlim=latlim)
  m = axs[1].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y, cmap=cmap, norm=norm, \
                                 shading='auto', extend='neither')
  axs[1].colorbar(m,loc='b',ticks=500.)
  
  ps_dir=prfx+'/PNG/GLBb0.08/'
  file_ps='glb0.08_NA_ke-eke_'+year+'_'+expt+'.png'
  print(file_ps)
  fig.savefig(ps_dir+file_ps,dpi=200,\
              facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps
  plot.show() ## interactive (block=False), otherwise ()

else:
  # 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
  fig, axs = plot.subplots(nrows=1,ncols=2,proj='mill',refheight='6cm',refwidth='10cm',tight=True)
  
  # KE
  #GS
  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 '+expt+' '+year,lonlim=lonlim, latlim=latlim)
  m = axs[0].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y0, cmap=cmap, norm=norm0, \
                                 shading='auto', extend='neither')
  axs[0].colorbar(m,loc='r',ticks=500.)
  
  
  # EKEM
  #GS
  axs[1].format(latlines=5, lonlines=5, labels=True,land=True,coast=True,landcolor='w',\
                title='EKE (cm2/s2) GLBb0.08 MOM6 '+expt+' '+year,lonlim=lonlim, latlim=latlim)
  m = axs[1].pcolormesh(lon_model[0:jdm-1,:], lat_model[0:jdm-1,:], y, cmap=cmap, norm=norm, \
                                 shading='auto', extend='neither')
  axs[1].colorbar(m,loc='r',ticks=500.)
  
  ps_dir=prfx+'/PNG/GLBb0.08/'
  file_ps='glb0.08_GS_ke-eke_'+year+'_'+expt+'.png'
  print(file_ps)
  fig.savefig(ps_dir+file_ps,dpi=200,\
              facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps
  plot.show() ## interactive (block=False), otherwise ()