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

import myenv as my

# get the grid
idm=4500; jdm=3298
kdm=34
iot='/nexsan/people/abozec/GLBb0.08/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

# data path
io='/Net/gleam/abozec/HYCOM/GLBb0.08/XIAOBIAO/expt_03.9/data/'
# get depth
dsd=my.nc.Dataset(io+'039_archm.1958_3zts.nc')
depth=dsd['Depth'][:]
print(depth)

# get basin masks
iom='/home/abozec/IDL/WORK/GLBb0.08/'
filem='mask_a-p-i_glbb08.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_p=my.np.reshape(tmp[1*(layer_size):2*layer_size],(jdm,idm)) ## Pacific
mask_i=my.np.reshape(tmp[2*(layer_size):3*layer_size],(jdm,idm)) ## Indian
mask_g=my.np.ones([jdm,idm])                                     ## Global

# choose basin to do
mask=mask_g
basin='Global Ocean'

# get length of the time series
tdm=12

# definition of array
tem=my.np.zeros([tdm,kdm])
sal=my.np.zeros([tdm,kdm])
t=0

for year in my.np.arange(tdm)+1958:
   # get file
   y='{:04d}'.format(year)
   print(y)
   file='039_archm.'+y+'_3zts.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])
   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.sum(pscx*pscy*mask3d[k,:,:])

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


# get the drift
drift_t=my.np.zeros([kdm,tdm])
drift_s=my.np.zeros([kdm,tdm])
for t in my.np.arange(tdm):
    drift_t[:,t]=tem[t,:]-tem[0,:]
    drift_s[:,t]=sal[t,:]-sal[0,:]

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

# plot drift
timey=my.np.arange(tdm)+1958
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(timey,depth,drift_t, cmap=cmap64,levels=ltemp)
axs[0].format(title='GLBb0.08 039 temp drift '+basin,ylim=(1000,0))

axs[1].contourf(timey,depth,drift_s, cmap=cmap64,levels=lsaln)
axs[1].format(title='GLBb0.08 039 saln drift '+basin,ylim=(1000,0))

m=axs[2].contourf(timey,depth,drift_t, cmap=cmap64,levels=ltemp)
axs[2].format(ylim=(6000,1000))
axs[2].colorbar(m, loc='b',ticks=0.5)

n=axs[3].contourf(timey,depth,drift_s, cmap=cmap64,levels=lsaln)
axs[3].format(ylim=(6000,1000))
axs[3].colorbar(n, loc='b',ticks=0.05)

# ps_dir='/p/home/abozec/PYTHON/PNG/'
# file_ps='GLBb0.08_ts_zt_1958-2018.pdf'
# print(file_ps)
# fig.savefig(ps_dir+file_ps,dpi=150,\
#             facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps
#

my.plot.show()