#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Nov 22 22:51:22 EST 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = sys.argv[1]
if 'GMST' in ifile:
    isGMST = True
else:
    isGMST = False
ds = xr.open_dataset(ifile)
#convert calendar to the standard gregorian
time_new = [datetime.datetime(t.year, t.month, t.day) for t in ds.time.values]
ds['time'] = time_new
ofile = ifile.replace('.nc', '.gregorian.nc')
#rename 'en' to 'ens'
ds = ds.rename(en='ens')
# transpose to put the time dimension first
if 'transpose' in sys.argv and not isGMST:
    ofile = ofile.replace('.nc', '.transpose.nc')
    ds = ds.transpose()
#ensemble-mean annual-mean
if isGMST:
    ofile = ofile.replace('.nc', '.ensAnnualMean.nc')
    time = ds.time.values[5::12]#Jul month every year
    ds = ds.mean('ens').groupby('time.year').mean('time') \
        .rename(year='time').assign_coords(time=time) #ens-annual-mean and assign new time
#if Jul-Jun mean
if 'Jul-Jun2years' in sys.argv:
    ofile = ofile.replace('.nc', '.Jul-Jun2years.nc')
    ds = ds.isel(time=slice(6,None)).rolling(time=24).mean() \
        .isel(time=slice(23,None,12))
elif 'Jul-Jun1year' in sys.argv:
    ofile = ofile.replace('.nc', '.Jul-Jun1year.nc')
    ds = ds.isel(time=slice(6,None)).rolling(time=12).mean() \
        .isel(time=slice(11,None,12))
#split ensemble members
if 'split_ens' in sys.argv:
    for en in ds.ens.values:
        ofile_en = ofile.replace('.nc', f'.en{en:02d}.nc')
        ds.sel(ens=en).drop('ens').to_netcdf(ofile_en)
        print('[saved]:', ofile_en)
else:
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
    
 
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()