#!python3
## get time series of maximum AMOC and GMOC at different latitude
import sys
prfx='/home/abozec/PYTHON/'
sys.path.append(prfx)

import myenv as my

## input directoryy1=2011 ; y2=2019
hycom_version=2.2
y1=1958 ; y2=2018
nt=(y2-y1+1) ## number of years of the time series
ex=701
PS=0     ## save a figure
file_nc= ## output an Netcdf file

## get year and expt formatted
year1='{:04d}'.format(y1)
year2='{:04d}'.format(y2)
expt='{:03d}'.format(ex)
expt1='{:04.1f}'.format(ex*0.1)
print('Expt:'+expt,'Expt1:'+expt1)
print(F'Number of years: {nt}')

## input directory and file
io='/Net/gleam/abozec/HYCOM/JRA-55/expt_'+expt1+'/data/'
if (file_nc == 1 ):
  fnc_out=io+'time_series_stmf_45N-41N-26N-34S_glbt_'+expt+'_'+year1+'_'+year2+'.nc'

# declaration of arrays for the different time series
stmf_max45 = my.np.zeros([nt])
stmf_max41 = my.np.zeros([nt])
stmf_max26 = my.np.zeros([nt])
stmf_max34a = my.np.zeros([nt])
stmf_min34g = my.np.zeros([nt])

n=0
for y in my.np.arange(nt)+y1:
    year='{:04d}'.format(y)
    file = expt+'_archm.'+year+'_strmf.nc'
    ds_stmf=my.nc.Dataset(io+file)
    stmf_atl=ds_stmf['strmf_atl'][0,:,:]
    stmf_glb=ds_stmf['strmf_glb'][0,:,:]
    lati=ds_stmf['Latitude'][:]
    nj=lati.size
    lat=lati.data
    ## find the min dy
    dy=my.np.nanmin([ my.np.abs( lat[0:nj-1] - my.np.roll(lat[0:nj-1],1,axis=0))] )

    # find indices for different latitude
    j45 =my.np.nanmax( my.np.where( (lat >  45.0-2.0*dy) & (lat <  45.0+2.0*dy) ) )
    j41 =my.np.nanmax( my.np.where( (lat >  41.0-2.0*dy) & (lat <  41.0+2.0*dy) ) )
    j26 =my.np.nanmax( my.np.where( (lat >  26.0-2.0*dy) & (lat <  26.0+2.0*dy) ) )
    j34s=my.np.nanmax( my.np.where( (lat > -34.0-2.0*dy) & (lat < -34.0+2.0*dy) ) )
    #print(F'j45: {j45}, j41: {j41}, j26: {j26}, j34s: {j34s}')

    # max at 45N
    stmf_max45[n] = my.np.nanmax(stmf_atl[:,j45])
    # max at 41N
    stmf_max41[n] = my.np.nanmax(stmf_atl[:,j41])
    # max at 26N
    stmf_max26[n] = my.np.nanmax(stmf_atl[:,j26])
    # max at 34S
    stmf_max34a[n] = my.np.nanmax(stmf_atl[:,j34s])
    stmf_min34g[n] = my.np.nanmin(stmf_glb[:,j34s])

    print(F'Year: {year}, strmf at 26N: {stmf_max26[n]}')
    n=n+1


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

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

    stmf45T = dso.createVariable('stmf_45N', 'f4', ('time',))
    stmf45T.standard_name = "max_strmf_45N"
    stmf45T.frequency = "year"
    stmf45T.units = "Sv"
    stmf45T[:] = stmf_max45[:]

    stmf41T = dso.createVariable('stmf_41N', 'f4', ('time',))
    stmf41T.standard_name = "max_strmf_41N"
    stmf41T.frequency = "year"
    stmf41T.units = "Sv"
    stmf41T[:] = stmf_max41[:]

    stmf26T = dso.createVariable('stmf_26N', 'f4', ('time',))
    stmf26T.standard_name = "max_strmf_26N"
    stmf26T.frequency = "year"
    stmf26T.units = "Sv"
    stmf26T[:] = stmf_max26[:]

    stmf34aT = dso.createVariable('stmf_atl_34S', 'f4', ('time',))
    stmf34aT.standard_name = "max_strmf_atl_34S"
    stmf34aT.frequency = "year"
    stmf34aT.units = "Sv"
    stmf34aT[:] = stmf_max34a[:]

    stmf34gT = dso.createVariable('stmf_glb_34S', 'f4', ('time',))
    stmf34gT.standard_name = "min_strmf_glb_34S"
    stmf34gT.frequency = "year"
    stmf34gT.units = "Sv"
    stmf34gT[:] = stmf_min34g[:]

    ## global attributes
    dso.source = "HYCOM "+'{:3.1f}'.format(hycom_version)
    dso.title = "Time series of yearly maximum vert. strmf. -GLBt0.72 config-"
    dso.institution = "FSU-COAPS"
    dso.history = "python"

    # close Netcdf
    dso.close()


## plot results
time=my.np.arange(nt) + y1

fig, ax = my.plot.subplots(ncols=1,nrows=3, aspect=3,axwidth='12cm',share=0, span=False)
ax[0].plot(time,stmf_max26, lw=2,legend='ul',labels=['26N'])
ax[0].plot(time,my.np.zeros([nt])+10.,lw=1,color='k',linestyle=':')
ax[0].format(title='AMOC transport at 26N', xlabel='Time (year)', ylabel='Vert. Stmf. (Sv)',\
             xlim=(y1,y2),ylim=(0,30),xlocator=10,ylocator=5)

ax[1].plot(time,stmf_max34a, lw=2,legend='ul',labels=['34S'])
ax[1].plot(time,my.np.zeros([nt])+10.,lw=1,color='k',linestyle=':')
ax[1].format(title='AMOC transport at 34S', xlabel='Time (year)', ylabel='Vert. Stmf. (Sv)', \
             xlim=(y1,y2),ylim=(0,30),xlocator=10,ylocator=5)

ax[2].plot(time,stmf_min34g, lw=2,legend='ul',labels=['34S'])
ax[2].plot(time,my.np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[2].format(title='GMOC transport at 34S', xlabel='Time (year)', ylabel='Vert. Stmf. (Sv)',\
             xlim=(y1,y2),ylim=(-50,0),xlocator=10,ylocator=10)


if (PS == 1):
   png_dir=io+'/PNG/'
   file_png='time_stmf_glbt072_'+expt+'_'+year1+'_'+year2+'.png'
   print(file_png)
   fig.savefig(png_dir+file_png,dpi=150,\
            facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps


my.plot.show()