#!/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]

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

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

## domain
idm=4500 ; jdm=3298
nk=80
PNG=1
#png_dir=prfx+'/PNG/ice_ocean_SIS2_01.0/'
png_dir=prfx+'/PNG/GLBb0.08/'
print('Starting here...')


## get path
y1=2010 ; y2=2014
year1='{:04d}'.format(y1) ; year2='{:04d}'.format(y2)
tdm=y2-y1+1

# get Netcdf
FILE_NC = 1
## output repository
expt1='01.0' ; expt='010'
#io_out=prfx+'../MOM6-expt/ice_ocean_SIS2_'+expt1+'/COMBINED/'
#fnc_out = io_out+'time_series_strmf_max-min_'+expt+'_glbbm0.25_om4_'+year1+'_'+year2+'.nc'

io_out=prfx+'../MOM6-expt/GLBb0.08/RLX_SSH/YEARLY/netcdf/'
fnc_out = io_out+'time_series_strmf_max-min_glb8_om4.1+++_'+year1+'_'+year2+'.nc'

#io=prfx+'../MOM6-expt/ice_ocean_SIS2_01.0/COMBINED/'
io=prfx+'../MOM6-expt/GLBb0.08/RLX_SSH/YEARLY/netcdf/'
psiPlot=np.zeros([tdm,nk,jdm])
psiPlotg=np.zeros([tdm,nk,jdm])
for y in np.arange(tdm)+y1:
  year='{:04d}'.format(y)
  file='archMNA.'+year+'_strmf.nc'
  ds=nc.Dataset(io+file)
  
  # get vmo
  strmf_atl=ds['strmf_atl'][0,:,:]
  strmf_atl[strmf_atl > 1e10] = 0. 
  strmf_glb=ds['strmf_glb'][0,:,:]
  strmf_glb[strmf_glb > 1e10] = 0.
  lat=ds['Latitude'][:]
  dep=ds['Depth'][:]
  #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[:] 
  
  
  # atlantic and global streamfunction
  psiPlot[y-y1,:,:] = strmf_atl[:,:]
  psiPlotg[y-y1,:,:] = strmf_glb[:,:]

# get latitude of the profile#
# max at 26N and max at 34S
max_26N=np.zeros([tdm])
max_34S=np.zeros([tdm])
min_34S=np.zeros([tdm])
for t in np.arange(tdm):
  zz,ii,psi=findExtrema(lata, depth, psiPlot[t,:,:], min_lat=26.5, max_lat=27., min_depth=250.) # RAPID
  max_26N[t]=psi
  zz1,ii1,psi=findExtrema(lata, depth, psiPlot[t,:,:], min_lat=-34.2, max_lat=-33.8, min_depth=250.) # RAPID
  max_34S[t]=psi
  zz2,ii2,psi=findExtremaMin(lata, depth, psiPlotg[t,:,:], min_lat=-34.2, max_lat=-33.8, min_depth=250.) # RAPID
  min_34S[t]=psi

print(max_26N[:])
print(max_34S[:])
print(min_34S[:])
print(ii,ii1,ii2)
print(lat[ii],lat[ii1],lat[ii2])

if (FILE_NC == 1):
  ### save data in file
  dso = nc.Dataset(fnc_out, 'w', format='NETCDF4')
  MT_dim = dso.createDimension('time', None)
  z_dim = dso.createDimension('depth', nk)

  MT = dso.createVariable('time', 'f4', ('time',))
  MT.long_name = "time"
  MT.units =  F"year since {year1}-01-01 00:00:00"
  MT.calendar = "gregorian"
  MT.axis = "T"
  MT[:]=np.arange(tdm)

  dep = dso.createVariable('depth', 'f4', ('depth',))
  dep.standard_name="level depth"
  dep.units = "m"
  dep[:] =  depth[:,ii]

  amoc = dso.createVariable('amoc_26N', 'f4', ('time','depth',))
  amoc.standard_name = "Atlantic Meridional Ocean Circulation at 26.5N"
  amoc.frequency = "year"
  amoc.units = "Sv"
  amoc[:,:] = psiPlot[:,:,ii]

  mx26N = dso.createVariable('atl_max_26N', 'f4', ('time',))
  mx26N.standard_name = "Maximum Atlantic streamfunction at 26.5N"
  mx26N.frequency = "year"
  mx26N.units = "Sv"
  mx26N[:] = max_26N[:]

  mx34S = dso.createVariable('atl_max_34S', 'f4', ('time',))
  mx34S.standard_name = "Maximum Atlantic streamfunction at 34S"
  mx34S.frequency = "year"
  mx34S.units = "Sv"
  mx34S[:] = max_34S[:]

  mn34S = dso.createVariable('glb_min_34S', 'f4', ('time',))
  mn34S.standard_name = "Minimum Global streamfunction at 34S"
  mn34S.frequency = "year"
  mn34S.units = "Sv"
  mn34S[:] = min_34S[:]

  ## global attributes
  dso.source = "MOM6"
  dso.title = "Time series of yearly Streamfunction -GLBb0.08 config-"
  dso.institution = "FSU-COAPS"
  dso.history = "python"

  # close Netcdf
  dso.close()


# 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.cycle='tab20'
pplot.rc.cycle='tab10'

pplot.rc.update({'fontname': 'Source Sans Pro', 'fontsize': 12})
array=[[1,2],[1,3],[1,4]]
time=np.arange(tdm)+y1

#fig, axs = pplot.subplots(nrows=1,refwidth=5,refheight=8)
fig, axs = pplot.subplots(array,refwidth=5,refheight=8,span=False)

# 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].plot(strmf_atl[:,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) GLBb0.08 Y:'+year1+'-'+year2,\
              ylim=(6500,0),xlim=(-10., 25.),\
              xlabel='AMOC (Sv)',ylabel='Depth (m)',\
              suptitle='')

# time series of streamfunctions
axs[1].plot(time,max_26N, lw=2)
axs[1].plot(time,np.zeros([tdm])+10.0,lw=0.5,ls='--',color='k')
axs[1].format(title='AMOC Maximum at 26.5N (Sv) ', xlabel='Time (year)', \
          ylabel='Streamfunction (Sv)',ylim=( 0.0, 30))

axs[2].plot(time,max_34S, lw=2)
axs[2].plot(time,np.zeros([tdm])+10.0,lw=0.5,ls='--',color='k')
axs[2].format(title='AMOC Maximum at 34S', xlabel='Time (year)', \
          ylabel='Streamfunction (Sv)',ylim=(0.0, 30.))

axs[3].plot(time,min_34S,lw=2)
axs[3].format(title='GMOC Minimum at 34S', xlabel='Time (year)', \
          ylabel='Streamfunction (Sv)',ylim=(-45.0,0))




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


pplot.show()