#!python3
## get time series of T,S above or below a certain depth from GLBm0.25
import sys
#prfx='/discover/nobackup/projects/gmao/cal_ocn/abozec1/PYTHON/'
prfx='/Users/abozec/PYTHON/'
sys.path.append(prfx)

import os
import proplot as plot
import netCDF4 as nc
import xarray as xr
import numpy as np
from myutilities.grid_mom6 import get_MOM6grid
from myutilities.tvplus import tvplus
from myutilities.get_cmap import get_cmaplct64
from dask.distributed import Client

## get experiment and time period
y1=1980 ; y2=1980
ex=10

## get year and expt formatted
ny=y2-y1+1 ; nt=(y2-y1+1)*12 ## get monthly files
year1='{:04d}'.format(y1) ; year2='{:04d}'.format(y2)
expt='{:03d}'.format(ex) ; expt1='{:04.1f}'.format(ex*0.1)
print('Expt:'+expt,'Expt1:'+expt1)

# depth over whic the calculate the tracers
depv1=700 # depth over which to calculate T and S
depv2=2000 # depth over which to calculate T and S
deps1='{:04d}'.format(int(depv1))
deps2='{:04d}'.format(int(depv2))
deps=deps1+'_'+deps2

# get Netcdf
FILE_NC = 0
## output repository
io_out=prfx+'../MOM6-expt/ice_ocean_SIS2_'+expt1+'/COMBINED/'
fnc_out = io_out+'time_series_TS-'+deps+'_'+expt+'_glbbm0.25_'+year1+'_'+year2+'.nc'

# Figure in PNG
PNG = 0
png_dir=prfx+'/PNG/ice_ocean_SIS2_'+expt1+'/'

# DEBUG -make sure the calculation is accurate for a new config
DEBUG = 0

# calculation of fields or reading of fields
CAL = 0

if (CAL == 1):
  ## get grid
  iog=prfx+'../Downloads/NASA/ice_ocean_SIS2_01.0/'
  #iog=prfx+'../MOM6-expt/ice_ocean_SIS2_'+expt1+'/'
  file='sea_ice_geometry.nc'
  idm=1440 ; jdm=1080 ; kdm=75
  dgrid=nc.Dataset(iog+file)
  #dgrid=xr.open_dataset(iog+file)
  #dgrid=xr.open_dataset(iog+file,chunks={'lath': 500, 'lonh': 500})
  dx=dgrid['dxT'][:]
  dy=dgrid['dyT'][:]
  mask=dgrid['wet'][:]
  area=dx*dy*mask
  area3d=np.tile(area,(kdm,1,1))
  print(area3d.shape)
  #area_tot=np.nansum(area[:,:])
  # get mask2d into a 3d
  ### mask3d = np.tile(mask,(kdm,1,1))
  
  ## input repository
  #io=prfx+'/../MOM6-expt/ice_ocean_SIS2_'+expt1+'/COMBINED/'
  io=prfx+'../Downloads/NASA/ice_ocean_SIS2_01.0/'
  
  ## define arrays
  temp_up1=np.zeros([nt]) ; temp_up2=np.zeros([nt])
  saln_up1=np.zeros([nt]) ; saln_up2=np.zeros([nt])
  temp_glb=np.zeros([nt]) ; saln_glb=np.zeros([nt])
  
  if (DEBUG == 1):
      # global TS
      temp_tot1=np.zeros([nt]) ; temp_tot2=np.zeros([nt])
      saln_tot1=np.zeros([nt]) ; saln_tot2=np.zeros([nt])
      # below ref dephs T and S
      temp_down1=np.zeros([nt]) ; temp_down2=np.zeros([nt])
      saln_down1=np.zeros([nt]) ; saln_down2=np.zeros([nt])
  
  ## fields to read in archive
  fields=['potT','salt','h']
  
  t=0
  year='1980'
  mon='01'
  for y in np.arange(ny)+y1:
    year='{:04d}'.format(y)
      for m in np.arange(1):
          mon='{:02d}'.format(m+1)
          file='ocean_month_'+year+'_'+mon+'.nc'
          print(io+file)
          mom6_field= nc.Dataset(io+file)
  
          ## get thickness
          tmp=mom6_field[fields[2]][0,:,0:jdm,0:idm]
          h=tmp.data ; h[h > 1e10] = np.nan
          h[h<0.01]=0.0
          print('h done !')
          ## get the interface depth
          dep = np.empty([kdm+1,jdm,idm])
          for k in np.arange(kdm)+1:
            dep[k,:,:] = dep[k-1,:,:] + h[k-1,:,:]
  
          dep[np.isnan(dep) == True] = 1e20
          print('dep done !')
  
          ## T
          tem=mom6_field[fields[0]][0,:,0:jdm,0:idm]
          #tem=tmp.values
          print('tem done !')
          ## S
          sal=mom6_field[fields[1]][0,:,0:jdm,0:idm]
          #sal=tmp.values
          print('sal done !')
  
          # get 3d volume
          vol=h*area3d
  
          # mask everything below depv1
          ind=np.ones([kdm,jdm,idm])
          mask1=np.where(dep[1:kdm+1,:,:] <= depv1,ind,0)
          mask2=np.where(dep[1:kdm+1,:,:] <= depv2,ind,0)
          # get indices of 700m layer
          kd1=np.array(np.sum(mask1,axis=0),dtype='int')
          kd2=np.array(np.sum(mask2,axis=0),dtype='int')
          depu1=np.zeros([jdm,idm])
          depu2=np.zeros([jdm,idm])
          depu1[kd1 == 75] = np.nan
          depu2[kd2 == 75] = np.nan
          kd1[kd1 == 75] = 0
          kd2[kd2 == 75] = 0
  
          # volume in upper layers abive depv
          volu1=np.nansum(vol*mask1)
          #print(F'vol1: {volu1}')
          volu2=np.nansum(vol*mask2)
          #print(F'vol2: {volu2}')
          # get the upper interface depth of the depv1 layers
          for j in np.arange(jdm):
              for i in np.arange(idm):
                 depu1[j,i]=dep[kd1[j,i],j,i]
                 depu2[j,i]=dep[kd2[j,i],j,i]
  
          depu1[depu1 > 1e10]=np.nan
          depu1[dep[kdm,:,:] < depv1]=np.nan
          depu2[depu2 > 1e10]=np.nan
          depu2[dep[kdm,:,:] < depv2]=np.nan
          # get the upper volume of the layer including depv1
          vol_lay1=(depv1-depu1)*area
          vol_lay2=(depv2-depu2)*area
          #print(F'uplay1: {np.nansum(vol_lay1)}')
          #print(F'uplay2: {np.nansum(vol_lay2)}')
          # get total volume above depv1
          vol_tot1=volu1+np.nansum(vol_lay1)
          vol_tot2=volu2+np.nansum(vol_lay2)
          #print(F'vol_tot1: {vol_tot1}')
          #print(F'vol_tot2: {vol_tot2}')
  
          # get temperature and salinity above depv1 and depv2
          tem1=np.nansum(tem*vol*mask1[:,:,:])/volu1
          tem2=np.nansum(tem*vol*mask2[:,:,:])/volu2
          sal1=np.nansum(sal*vol*mask1[:,:,:])/volu1
          sal2=np.nansum(sal*vol*mask2[:,:,:])/volu2
          temi1=np.zeros([jdm,idm]) ; temi2=np.zeros([jdm,idm])
          sali1=np.zeros([jdm,idm]) ; sali2=np.zeros([jdm,idm])
          for j in np.arange(jdm):
              for i in np.arange(idm):
                 temi1[j,i]=tem[kd1[j,i],j,i]
                 temi2[j,i]=tem[kd2[j,i],j,i]
                 sali1[j,i]=sal[kd1[j,i],j,i]
                 sali2[j,i]=sal[kd2[j,i],j,i]
  
          vol_tem1=np.nansum(temi1*vol_lay1)/np.nansum(vol_lay1)
          vol_tem2=np.nansum(temi2*vol_lay2)/np.nansum(vol_lay2)
          temp_up1[t]=(tem1*volu1 + vol_tem1*np.nansum(vol_lay1))/vol_tot1
          temp_up2[t]=(tem2*volu2 + vol_tem2*np.nansum(vol_lay2))/vol_tot2
  
          vol_sal1=np.nansum(sali1*vol_lay1)/np.nansum(vol_lay1)
          vol_sal2=np.nansum(sali2*vol_lay2)/np.nansum(vol_lay2)
          saln_up1[t]=(sal1*volu1 + vol_sal1*np.nansum(vol_lay1))/vol_tot1
          saln_up2[t]=(sal2*volu2 + vol_sal2*np.nansum(vol_lay2))/vol_tot2
          ## print results
          print(F't: {t}')
          print(F'T_up1: {temp_up1[t]}, S_up1: {saln_up1[t]}')
          print(F'T_up2: {temp_up2[t]}, S_up2: {saln_up2[t]}')
  
          if (DEBUG == 1):
              #T,S, vol and depth below depv
              depd1[n]=np.nanmax([depv1,dep1d[kd1,n]])
              if (kd1 <= kdm-1):
                vold1[kd1:kdm,n]=vol1d[kd1:kdm,n]
              else:
                vold1[kd1:kdm,n]=0.0
     
              depd1[depd1 > 1e10]=np.nan
              depd2[depd2 > 1e10]=np.nan
              depd2[dep[kdm,:,:] < depv2]=np.nan
              depd1[dep[kdm,:,:] < depv1]=np.nan
     
          ## TS below depv
              # volume below depv layer
              do_vol=np.nansum(vold1)
              # volume of the bottom of the layer containing depv
              depv_vol=(depd1 - depv1)*area1d # 2-D
              depv_dow_vol=np.nansum(depv_vol)
              # total volume above depv
              do_vol_tot1=do_vol+depv_dow_vol
     
              # temperature below depv
              do_tem=np.nansum(temd1*vold1)/do_vol
              depv_tem=np.nansum(temi1*depv_vol)/depv_dow_vol
              temp_down1[t]=(do_tem*do_vol + depv_tem*depv_dow_vol)/do_vol_tot1
     
              # salinity below depv
              do_sal=np.nansum(sald1*vold1)/do_vol
              depv_sal=np.nansum(sali1*depv_vol)/depv_dow_vol
              saln_down1[t]=(do_sal*do_vol + depv_sal*depv_dow_vol)/do_vol_tot1
     
              ## sanity check with total volume T and S
              temp_tot1[t]=(temp_up1[t]*up_vol_tot1+temp_down1[t]*do_vol_tot1)/(up_vol_tot1+do_vol_tot1)
              saln_tot1[t]=(saln_up1[t]*up_vol_tot1+saln_down1[t]*do_vol_tot1)/(up_vol_tot1+do_vol_tot1)
              temp_tot2[t]=(temp_up2[t]*up_vol_tot2+temp_down2[t]*do_vol_tot2)/(up_vol_tot2+do_vol_tot2)
              saln_tot2[t]=(saln_up2[t]*up_vol_tot2+saln_down2[t]*do_vol_tot2)/(up_vol_tot2+do_vol_tot2)
     
              print(F'T_tot1: {temp_tot1[t]}, S_tot1: {saln_tot1[t]}')
              print(F'T_tot2: {temp_tot2[t]}, S_tot2: {saln_tot2[t]}')
     
              print(F'vol_glb: {vol1_tot}')
              print(F'vol_tot: {up_vol_tot+do_vol_tot}')
              print("T_tot :{0:9.5f}, S_tot :{1:9.5f}".format(temp_tot[t],saln_tot[t]))
     
              ## global T and S
              temp_glb[t]=np.nansum(tem1d*vol1d)/vol1_tot
              saln_glb[t]=np.nansum(sal1d*vol1d)/vol1_tot
              print("T_glb :{0:9.5f}, S_glb :{1:9.5f}".format(temp_glb[t],saln_glb[t]))
  
          t=t+1
  
  if (FILE_NC == 1):
    ## save data in file
    dso = 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"months since {year1}-01-01 00:00:00"
    MT.calendar = "gregorian"
    MT.axis = "T"
    MT[:]=np.arange(nt)
    
    temp = dso.createVariable('thetaoga', 'f4', ('time',))
    temp.standard_name = "average_sea_water_potential_temperature"
    temp.frequency = "year"
    temp.units = "degC"
    temp[:] = temp_glb[:]
    
    saln = dso.createVariable('soga', 'f4', ('time',))
    saln.standard_name = "average_sea_water_salinity"
    saln.frequency = "year"
    saln.units = ""
    saln[:] = saln_glb[:]
    
    tempT = dso.createVariable('thetaoga_top'+deps1, 'f4', ('time',))
    tempT.standard_name = "0-"+deps1+"m_average_sea_water_potential_temperature"
    tempT.frequency = "year"
    tempT.units = "degC"
    tempT[:] = temp_up1[:]
    
    salnT = dso.createVariable('soga_top'+deps1, 'f4', ('time',))
    salnT.standard_name = "0-"+deps1+"m_average_sea_water_salinity"
    salnT.frequency = "year"
    salnT.units = ""
    salnT[:] = saln_up1[:]
    
    tempB = dso.createVariable('thetaoga_top'+deps2, 'f4', ('time',))
    tempB.standard_name = "0-"+deps2+"m_average_sea_water_potential_temperature"
    tempB.frequency = "year"
    tempB.units = "degC"
    tempB[:] = temp_up2[:]
    
    salnB = dso.createVariable('soga_top'+deps2, 'f4', ('time',))
    salnB.standard_name = "0-"+deps2+"m_average_sea_water_salinity"
    salnB.frequency = "year"
    salnB.units = ""
    salnB[:] = saln_up2[:]
    
    ## global attributes
    dso.source = "MOM6-SIS2"
    dso.title = "Time series of monthly tracers -GLBm0.25 config-"
    dso.institution = "FSU-COAPS"
    dso.history = "python"
    
    # close Netcdf
    dso.close()

else:
  dso=nc.Dataset(fnc_out)
  temp_glb=dso['thetaoga'][:]
  saln_glb=dso['soga'][:]
  temp_up1=dso['thetaoga_top700'][:]
  saln_up1=dso['soga_top700'][:]
  temp_up2=dso['thetaoga_top2000'][:]
  saln_up2=dso['soga_top2000'][:]
  nt=temp_up1.size

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

fig, ax = plot.subplots(ncols=2,nrows=3, refwidth='10cm',share=0, span=False)
# global Temperature and slainity anomaly
ax[0].plot(time,temp_glb-temp_glb[0], lw=2,legend='lr',labels=['Pot. Temp.'])
ax[0].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[0].format(title='Global pot. temperature drift', xlabel='Time (year)', \
          ylabel='Pot. Temp. (C)',ylim=(-0.1,0.1),xlocator=2,ylocator=0.05)

ax[1].plot(time,saln_glb-saln_glb[0], lw=2,legend='lr',labels=['Salinity'])
ax[1].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[1].format(title='Global salinity drift', xlabel='Time (year)', \
          ylabel='Salinity',ylim=(-0.02,0.02),xlocator=2,ylocator=0.005)

ax[2].plot(time,temp_up1-temp_up1[0], lw=2,legend='lr',labels=['Pot. Temp.'])
ax[2].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[2].format(title='Pot. Temperature above '+deps1+'m', xlabel='Time (year)', \
          ylabel='Pot. Temp. (C)',ylim=(-0.2,0.6),xlocator=2,ylocator=0.2)

ax[3].plot(time,saln_up1-saln_up1[0], lw=2,legend='lr',labels=['Salinity'])
ax[3].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[3].format(title='Salinity above '+deps1+'m ', xlabel='Time (year)', \
          ylabel='Salinity',ylim=(-0.1,0.1),xlocator=2,ylocator=0.05)

ax[4].plot(time,temp_up2-temp_up2[0], lw=2,legend='lr',labels=['Pot. Temp.'])
ax[4].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[4].format(title='Pot. Temperature above '+deps2+'m', xlabel='Time (year)', \
          ylabel='Pot. Temp. (C)',ylim=(-0.1,0.15),xlocator=2,ylocator=0.05)

ax[5].plot(time,saln_up2-saln_up2[0], lw=2,legend='lr',labels=['Salinity'])
ax[5].plot(time,np.zeros([nt]),lw=1,color='k',linestyle=':')
ax[5].format(title='Salinity above '+deps2+'m ', xlabel='Time (year)', \
          ylabel='Salinity',ylim=(-0.02,0.025),xlocator=2,ylocator=0.01)

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

plot.show()