#python3
# test xarray python with the GLbb0.08
import sys
prfx='/discover/nobackup/projects/gmao/cal_ocn/abozec1/PYTHON/'
sys.path.append(prfx)

import xarray as xr
import numpy as np
from myutilities.tvplus import tvplus
import time
from time import process_time
import dask
from dask.distributed import Client, progress

def weight_sum(var):
  io='/discover/nobackup/projects/gmao/cal_ocn/sakella/GLBb0.08/expt_10.6-02Oct2023/'
  year='2012'
  nd=2
  kdm=41
  idm=4500 ; jdm=3297
  #hs=np.zeros([kdm,jdm,idm])
  #var0=np.zeros([kdm,jdm,idm])
  hs=np.zeros([jdm,idm])
  var0=np.zeros([jdm,idm])
  t=0
  for d in np.arange(nd)+33:
      day='{:03d}'.format(d)
      print(day)
      start = process_time()
      ds=xr.open_dataset(io+'ocnm_'+year+'_'+day+'.nc')
      # get h
      hs  = hs   + ds.h[0,0,:,:].values
      # get potT
      var0= var0 + ds[var][0,0,:,:].values*ds.h[0,0,:,:].values
      t=t+1
      end = process_time()
      print('Process time for day',day,':',end-start)
  
  # get the average
  var0=var0/hs
  h=hs/t
  
  print('Min/Max var:',np.nanmin(var0[0,:,:]),np.nanmax(var0[0,:,:]))
  print('Min/Max h:',np.nanmin(h[0,:,:]),np.nanmax(h[0,:,:]))
  


io='/discover/nobackup/projects/gmao/cal_ocn/sakella/GLBb0.08/expt_10.6-02Oct2023/'
year=2012
nd=2
var='potT'
list=[io,year,nd]
time weight_sum(['potT'])

client=Client()
client

io='/discover/nobackup/projects/gmao/cal_ocn/sakella/GLBb0.08/expt_10.6-02Oct2023/'
ds=xr.open_mfdataset(io+'ocnm_2012_10*.nc')
havg=ds.h.mean(dim='Time')
Th=(ds.h*ds.potT).sum(dim='Time')
sumh=ds.h.sum(dim='Time')


#ds1=xr.open_mfdataset(io+'ocnm_2010_34*.nc',parallel=True) ## requires dask
#nt=ds1.Time.shape[0]
#sumh=ds1.h.sum(dim='Time').compute()
#sumh=sumh/nt
#
#Th=ds1.potT * ds1.h
#sumT=Th.sum(axis=0).compute()
#meanT=sumT/(sumh*nt)
#
#Sh=ds1.salt * ds1.h
#sumS=Sh.sum(axis=0).compute()
#meanS=sumS/(sumh*nt)