#!python3
# plot EKE of GLBb0.08
import sys
prfx='/home/abozec/PYTHON/'
sys.path.append(prfx)

import myenv as my

## Experiment number
ex=39
PS=0
expt='{:03d}'.format(ex)
expt1='{:04.1f}'.format(ex*0.1)
print('Expt:'+expt,'Expt1:'+expt1)

## input directory and file
io='/Net/gleam/abozec/HYCOM/GLBb0.08/XIAOBIAO/expt_'+expt1+'/data/'

## Model state
file=expt+'_archSDA.1993_2018_5d_eke_0005m.nc' # from get_eke_ncdf_glb08.csh script
factor=1. #  = 1 for 5m, =0.1 for 700m and 1000m
ds1=my.nc.Dataset(io+file)
tmp=ds1['kinetic_energy_per_mass'][0,0,:,:]
eke=tmp*1e4 # convert to cm2/s2
eke.data[eke.mask == True]=my.np.nan
print('Finished reading data')
# get grid
lon=ds1['Longitude'][:,:] ; lat=ds1['Latitude'][:,:]

## interpolate  every model point (ocean+mask) to 0.5 deg
eke_reg,grid_x,grid_y=my.interp2plot(eke,lon,lat,reso=0.5,ocean=False)
print('Finished interpolating data')

# plot the field
# get cmap
cmap64=my.mygc.get_cmaplct64()

# get the value levels to plot
vmin=0 ; vmax=1000.*factor
leke=my.np.linspace(vmin,vmax,21)
print(leke)
eke_reg[eke_reg > 1000*factor]=999.*factor ## saturate the color over 1000. for better comparison with paper

title1='GLBb0.08 EKE 5d avg 5m ('+expt+') Y: 1993-2018'

# plot the field
my.plot.rc.reso = 'med'  # use higher res for zoomed in geographic features
my.plot.rc.coast=False

proj = my.plot.Proj('cyl', basemap=True,lonlim=(0,360),latlim=(-80,90))
fig, axs = my.plot.subplots(ncols=1,nrows=1,proj=(proj),axwidth='6in',top='2em')
CS1=axs[0].contourf(grid_x,grid_y,eke_reg, cmap=cmap64, globe=True,\
                    levels=leke,extend='both')
axs[0].colorbar(CS1,loc='b',ticks=100.*factor)
axs[0].format(latlines=30., lonlines=30, labels=True,land=True,coast=True,\
              title=title1,landcolor='White')

if (PS ==1):
   ps_dir=io+'/PNG/'
   file_ps='EKE-5d_5m_'+expt+'_1993-2018.png'
   print(file_ps)
   fig.savefig(ps_dir+file_ps,dpi=150,\
            facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps

#my.plot.show(block=False)
my.plot.show()