#!python3
# plot yz  per basin of GLBt0.72
import sys
prfx='/home/abozec/PYTHON/'
sys.path.append(prfx)

import myenv as my

# get the grid
idm=500; jdm=382
kdm=60
iot='/nexsan/people/abozec/Net_yucatan/PACIFIC/HYCOM/hycom/GLBt0.72/topo/'
pscx=my.hio.sub_var2(iot+'regional.grid.a',idm,jdm,10) ## dx
pscy=my.hio.sub_var2(iot+'regional.grid.a',idm,jdm,11) ## dy
plat=my.hio.sub_var2(iot+'regional.grid.a',idm,jdm,2) ## Latitude
# get the latitude in a 1-dim array
zloc=my.np.where(plat[jdm-1,:] == my.np.max(plat[jdm-1,:]))
zlat=plat[:,zloc[0][0]] ## column of the highest latitude


# data path
io='/Net/gleam/abozec/HYCOM/JRA-55/expt_70.1/data/'
# get depth
dsd=my.nc.Dataset(io+'701_archm.1958_3zts-int.nc')
depth=dsd['Depth'][:]
print(depth)

# get basin masks
iom='/home/abozec/IDL/WORK/GLBt0.72/'
filem='mask_a-p-i_glbt.dat'
layer_size=idm*jdm
tmp=my.np.fromfile(iom+filem,dtype='<f')
mask_a=my.np.reshape(tmp[0:layer_size],(jdm,idm))                ## Atlantic
mask_a[0:110,:]=0. ## mask south of Africa's tip
mask_p=my.np.reshape(tmp[1*(layer_size):2*layer_size],(jdm,idm)) ## Pacific
mask_p[0:110,:]=0.
mask_i=my.np.reshape(tmp[2*(layer_size):3*layer_size],(jdm,idm)) ## Indian
mask_i[0:110,:]=0.
mask_g=my.np.ones([jdm,idm])                                     ## Global

# choose basin to do
mask=mask_i+mask_p
basin='Indo-Pacific Ocean'

# definition of array
tem=my.np.zeros([kdm,jdm])
sal=my.np.zeros([kdm,jdm])
tem_clim=my.np.zeros([kdm,jdm])
sal_clim=my.np.zeros([kdm,jdm])


# get file
y='1980-2018'
file='701_archm.'+y+'_3zts-int.nc'
# get the Netcdf variable
ds=my.nc.Dataset(io+file)
temp=ds['pot_temp'][0,:,:,:]
saln=ds['salinity'][0,:,:,:]

# get mask and area
area=my.np.zeros([kdm,jdm])
mask3d=my.np.ones([kdm,jdm,idm])
mask3d[temp.data > 1e5 ] = 0.
for k in my.np.arange(kdm):
  mask3d[k,:,:]=mask3d[k,:,:]*mask[:,:]
  area[k,:]=my.np.nansum(pscx*pscy*mask3d[k,:,:],axis=1)

#print(temp.data.shape)
#print(mask3d.shape)
# Zonal average
for k in my.np.arange(kdm):
  tem[k,:]=my.np.nansum(temp.data[k,:,:]*pscx*pscy*mask3d[k,:,:],axis=1)/area[k,:]
  sal[k,:]=my.np.nansum(saln.data[k,:,:]*pscx*pscy*mask3d[k,:,:],axis=1)/area[k,:]


# get clim
io_clim = '/Net/yucatan/abozec/PACIFIC/HYCOM/hycom/GLBt0.72/relax/relax_sig2_17t_gdem4_41l_wosigma/netcdf_60l/'
# get the Netcdf variable
ds=my.nc.Dataset(io_clim+'relax.tem.gdem4.nc')
temp_clim=ds['pot_temp'][0,:,:,:]
ds=my.nc.Dataset(io_clim+'relax.sal.gdem4.nc')
saln_clim=ds['salinity'][0,:,:,:]

# get mask and area
area=my.np.zeros([kdm,jdm])
mask3d=my.np.ones([kdm,jdm,idm])
mask3d[temp_clim.data > 1e5 ] = 0.
for k in my.np.arange(kdm):
  mask3d[k,:,:]=mask3d[k,:,:]*mask[:,:]
  area[k,:]=my.np.nansum(pscx*pscy*mask3d[k,:,:],axis=1)

#print(temp_clim.data.shape)
#print(mask3d.shape)
# zonal average
for k in my.np.arange(kdm):
  tem_clim[k,:]=my.np.nansum(temp_clim.data[k,:,:]*pscx*pscy*mask3d[k,:,:],axis=1)/area[k,:]
  sal_clim[k,:]=my.np.nansum(saln_clim.data[k,:,:]*pscx*pscy*mask3d[k,:,:],axis=1)/area[k,:]


# get the drift
drift_t=tem-tem_clim
drift_s=sal-sal_clim

# get cmap
cmap64=my.mygc.get_cmaplct64()
N=cmap64.colors.shape[0]
# colorbar bounds for Temperature
vmin=-2.5 ; vmax=2.5
ltemp=my.np.linspace(vmin,vmax,21)
# colorbar bounds for salinity
vmin=-0.5 ; vmax=0.5
lsaln=my.np.linspace(vmin,vmax,21)

# plot drift
fig, axs = my.plot.subplots(ncols=2,nrows=2, axwidth='10cm', span=False, \
           aspect=3,hspace=0,hratios=(1.5,1),ref=1)
axs[0].contourf(zlat,depth,drift_t, cmap=cmap64,levels=ltemp)
axs[0].contour(zlat,depth,drift_t,levels=ltemp,linewidths=0.2, colors='k')
axs[0].format(title='GLBt0.72 701 temp drift '+basin,ylim=(1500,0),xlim=(-80,80),ylocator=200)

axs[1].contourf(zlat,depth,drift_s, cmap=cmap64,levels=lsaln)
axs[1].contour(zlat,depth,drift_s,levels=lsaln,linewidths=0.2, colors='k')
axs[1].format(title='GLBt0.72 701 saln drift '+basin,ylim=(1500,0),xlim=(-80,80),ylocator=200)

m=axs[2].contourf(zlat,depth,drift_t, cmap=cmap64,levels=ltemp)
axs[2].contour(zlat,depth,drift_t,levels=ltemp,linewidths=0.2, colors='k')
axs[2].format(ylim=(6500,1500),ylocator=1000)
axs[2].colorbar(m, loc='b',ticks=0.5)

n=axs[3].contourf(zlat,depth,drift_s, cmap=cmap64,levels=lsaln)
axs[3].contour(zlat,depth,drift_s,levels=lsaln,linewidths=0.2, colors='k')
axs[3].format(ylim=(6500,1500),ylocator=1000)
axs[3].colorbar(n, loc='b',ticks=0.10)

# ps_dir=io+'PNG/'
# file_ps='GLBt0.72_tsbias_yz_global_1980-2018.pdf'
# print(file_ps)
# fig.savefig(ps_dir+file_ps,dpi=150,\
#             facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps
#

my.plot.show()