#!/python3
# plot bias of GLBm0.25
import sys
prfx='/discover/nobackup/projects/gmao/cal_ocn/abozec1/PYTHON/'
sys.path.append(prfx)

import os
import proplot as pplot
import netCDF4 as nc
import numpy as np
import xarray as xr
import pandas as pd
#from myutilities.tvplus import tvplus


####################################################################################################################
## Functions needed - adapted from mom6_tools from NCAR

def MOCpsi(vh, vmsk=None):
  """Sums 'vh' zonally and cumulatively in the vertical to yield an overturning stream function, psi(y,z)."""
  shape = list(vh.shape); shape[-3] += 1
  psi = np.zeros(shape[:-1])
  if len(shape)==3:
    for k in range(shape[-3]-1,0,-1):
      if vmsk is None: psi[k-1,:] = psi[k,:] - vh[k-1].sum(axis=-1)
      else: psi[k-1,:] = psi[k,:] - (vmsk*vh[k-1]).sum(axis=-1)
  else:
    for n in range(shape[0]):
      for k in range(shape[-3]-1,0,-1):
        if vmsk is None: psi[n,k-1,:] = psi[n,k,:] - vh[n,k-1].sum(axis=-1)
        else: psi[n,k-1,:] = psi[n,k,:] - (vmsk*vh[n,k-1]).sum(axis=-1)
  return psi

def findExtrema(y, z, psi, min_lat=-90., max_lat=90., min_depth=0., mult=1.):
  psiMax = mult*np.amax( mult * np.ma.array(psi)[(y>=min_lat) & (y<=max_lat) & (z>min_depth)] )
  idx = np.argmin(np.abs(psi-psiMax))
  (j,i) = np.unravel_index(idx, psi.shape)
  return j,i,psi[j,i]

####################################################################################################################

## domain
idm=1440 ; jdm=1081 
nk=36
PNG=1
png_dir=prfx+'/PNG/ice_ocean_SIS2_01.0/'
print('Starting here...')

# get basin masks
iom=prfx+'../DATA_OBS/GLBm0.25/'
filem='basin_codes.v20140629.nc'
dmsk=nc.Dataset(iom+filem)
tmp=dmsk['basin'][:,:]
basin=np.zeros([tmp.shape[0]+1,tmp.shape[1]])
basin[0:jdm-1,:]=tmp[:,:]
basin[jdm-1,:]=basin[jdm-2,:]
m = 0*basin; m[(basin==2) | (basin==4) | (basin==6) | (basin==7) | (basin==8)]=1

## get path
y1=1980 ; y2=1984
year1='{:04d}'.format(y1) ; year2='{:04d}'.format(y2)
tdm=y2-y1+1
io=prfx+'../MOM6-expt/ice_ocean_SIS2_01.0/COMBINED/'
psiPlot=np.zeros([tdm,nk,jdm])
for y in np.arange(tdm)+y1:
  year='{:04d}'.format(y)
  file='ocean_month_z_'+year+'.nc'
  ds=nc.Dataset(io+file)
  
  # get vmo
  conversion_factor=1e-9
  vmo=ds['vmo'][0,:,:,:]
  lat=ds['yq'][:]
  dep=ds['z_i'][:]
  #nk=dep.shape[0]
  
  # convert to 2d arrays
  lata=np.zeros([nk,jdm]) ; depth=np.zeros([nk,jdm])
  for k in np.arange(nk):lata[k,:]=lat[:]
  
  for j in np.arange(jdm):depth[:,j]=dep[:] 
  
  
  # global streamfunction
  psiPlot[y-y1,:,:] = MOCpsi(vmo[:,:,:],vmsk=m[:,:])*conversion_factor

# get latitude of the profile#
zz,ii,psi=findExtrema(lata, depth, psiPlot[0,:,:], min_lat=26.5, max_lat=27., min_depth=250.) # RAPID

# RAPid section
# RAPid section
io_obs='/discover/nobackup/projects/gmao/cal_ocn/abozec1/DATA_OBS/'
drapid=nc.Dataset(io_obs+'moc_vertical.nc')
stmf=drapid['stream_function_mar'][:]
depth_rapid=drapid['depth'][:]
nz=depth_rapid.shape[0]

# use panda Frame to do the shading of the standard deviation
data_stmf = pd.DataFrame(stmf.T, columns=depth_rapid)
means = data_stmf.mean(axis=0)
std   = data_stmf.std(axis=0)
shadedata = np.zeros([2,nz])
shadedata[0,:] = means + std
shadedata[1,:] = means - std

kw = dict(
    shadedata=shadedata,
    label='RAPID array', shadelabel='std-dev',
    color='gray6', barzorder=0, boxmarker=False,
)

print('Starting plotting...')
# plot profile
# format plot
pplot.rc.update({'fontname': 'Source Sans Pro', 'fontsize': 12})

fig, axs = pplot.subplots(nrows=1,refwidth=5,refheight=8)
# zeros line plots
axs[0].plot(np.zeros([depth_rapid.shape[0]]),depth_rapid,lw=1,linestyle='--',color='k')
# RAPID data
axs[0].linex(means, legend='ll', **kw)
# Model data
for t in np.arange(tdm):
  axs[0].plot(psiPlot[t,:,ii],depth[:,ii], lw=1.,legend='lr',\
        legend_kw={'ncols': 1,'labels':['AMOC 26.5N Y:'+'{:04d}'.format(t+y1)]})
axs[0].format(title='Atlantic Meridional Overturning Circulation at 26.5N (Sv) GLBm0.25 Y:'+year1+'-'+year2,\
              ylim=(6500,0),xlim=(-10., 25.),\
              xlabel='AMOC (Sv)',ylabel='Depth (m)')

if (PNG ==1):
   file_png='profile_amoc26N_glbm0.25_'+year1+'-'+year2+'.png'
   print(file_png)
   fig.savefig(png_dir+file_png,dpi=150,\
            facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps


pplot.show()