#!/python3
# plot bias of GLBm0.25
# to avoid attribute error due to version in proplot ...
from importlib.metadata import version
from packaging.version import parse
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
from hycom.info import read_field_names,read_field_grid_names
from hycom.io import read_hycom_fields, read_hycom_grid,sub_var2


####################################################################################################################
## 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=39
PNG=1
#png_dir=prfx+'/PNG/ice_ocean_SIS2_01.0/'
png_dir=prfx+'/PNG/GLBb0.08/'
print('Starting here...')

# get basin masks
iom=prfx+'../DATA_OBS/GLBb0.08/'
filem='basin_GLBb0.08_0-7_09m11.a'
tmp=sub_var2(iom+filem,idm,jdm,1,replace_to_nan=False)
basin=tmp
m = 0*basin; m[(basin==1)]=1

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

# get Netcdf
FILE_NC = 1
## output repository
expt1='10.8' ; expt='108'
#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/expt_'+expt1+'/MONTHLY_Z/'
fnc_out = io_out+'time_series_strmf_max-min_glb8_G16.9_'+expt+'_'+year1+'_'+year2+'.nc'

#io=prfx+'../MOM6-expt/ice_ocean_SIS2_01.0/COMBINED/'
io=prfx+'../MOM6-expt/GLBb0.08/expt_'+expt1+'/MONTHLY_Z/'
psiPlot=np.zeros([tdm,nk,jdm])
psiPlotg=np.zeros([tdm,nk,jdm])
for y in np.arange(tdm)+y1:
  year='{:04d}'.format(y)
  file='ocean_month_z_'+year+'.nc'
  print('Year:',year)
  ds=xr.open_dataset(io+file)
  #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[:] 
  

  # atlantic and global streamfunction
  psiPlot[y-y1,:,:] = MOCpsi(vmo[:,:,:],vmsk=m[:,:])*conversion_factor
  psiPlotg[y-y1,:,:] = MOCpsi(vmo[:,:,:])*conversion_factor

# 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[:])

# read the first 5 years from Netcdf
#dnc=nc.Dataset(io+'time_series_strmf_max-min_glb8_om4.1+++_2010_2014.nc')
#tmp=dnc['amoc_26N'][:,:]
#psiPlot[0:5,:,ii]=tmp[:,:]
#tmp=dnc['atl_max_26N'][:]
#max_26N[0:5]=tmp[:]
#tmp=dnc['atl_max_34S'][:]
#max_34S[0:5]=tmp[:]
#tmp=dnc['glb_min_34S'][:]
#min_34S[0:5]=tmp[:]
#
#print(max_26N[:])
#print(max_34S[:])
#print(min_34S[:])

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()


# previous years
# read the first 5 years from Netcdf
#nk0=80
#y1=2010 ; y2=2019
#psiPlot0=np.zeros([5,nk0])
#max0_26N=np.zeros([y2-y1+1])
#max0_34S=np.zeros([y2-y1+1])
#min0_34S=np.zeros([y2-y1+1])
#dnc=nc.Dataset(io+'time_series_strmf_max-min_glb8_om4.1+++_2010_2014.nc')
#tmp=dnc['amoc_26N'][:,:]
#psiPlot0[0:5,:]=tmp[:,:]
#tmp=dnc['atl_max_26N'][:]
#max0_26N[0:5]=tmp[:]
#tmp=dnc['atl_max_34S'][:]
#max0_34S[0:5]=tmp[:]
#tmp=dnc['glb_min_34S'][:]
#min0_34S[0:5]=tmp[:]
#
#max0_26N[5:10]=max_26N[:]
#max0_34S[5:10]=max_34S[:]
#min0_34S[5:10]=min_34S[:]
#
#depth0=dnc['depth'][:]
#
#print(max0_26N[:])
#print(max0_34S[:])
#print(min0_34S[:])





# RAPid section
io_obs='/discover/nobackup/projects/gmao/cal_ocn/abozec1/DATA_OBS/'
#drapid=nc.Dataset(io_obs+'moc_vertical.nc')
drapid=xr.open_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.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+2010)]})
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='G16.9 expt:'+expt1)



#for t in np.arange(5)+5:
#  axs[0].plot(psiPlot[t-5,:,ii],depth[:,ii], lw=1.,legend='lr',\
#        legend_kw={'ncols': 1,'labels':['AMOC 26.5N Y:'+'{:04d}'.format(t+2010)]})
#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='OM4.1+++')





# time series of streamfunctions
time=np.arange(tdm)+2010
print(max_26N.shape)
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_glb8_G16.9_'+expt+'_'+year1+'-'+year2+'.png'
#   file_png='profile_amoc26N_glb8_om4.1+++_2010-2019.png'
   print(file_png)
   fig.savefig(png_dir+file_png,dpi=150,\
            facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps


pplot.show()