#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Dec 18 23:32:43 EST 2019
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
from wy_vecchi_indices_test import sst_to_vecchi_indices
#print()
# README: test the impact when TAREA is not land masked

ifiles = 'LME_SST/b.e11.BLMTRC5CN.f19_g16.???.*.nc'
n_ens = 13
ens = range(1, n_ens+1)
ofile = f'LME_TC.fullForcingTest.{n_ens}ens.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
#sstName = 'SST'
#xName = 'nlon'
#yName = 'nlat'
#latName = 'TLAT'
#lonName = 'TLONG'
#areaName = 'TAREA'
#shiftTimeAxis = True #CESM data time axis is one month late and needs a backward shift
yearsRef = slice(1901, 2000) # reference years in calculation of sst anomaly for TC indices
#def sst_to_vecchi_indices(ifiles, ofile=None, sstName='SST', xName='nlon', yName='nlat', lonName='TLONG', latName='TLAT', areaName='TAREA', shiftTimeAxis=True, yearsRef=None):
     
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.datetime.now()
    print('[start]:', t0.strftime(tformat))
    
    dss = []
    for en in ens:
        print(f'en = {en:02d}/{n_ens:02d}')
        _ifiles = ifiles.replace('???', f'{en:03d}')
        ds = sst_to_vecchi_indices(_ifiles, yearsRef=yearsRef)
        dss.append(ds)

    # concat
    ds_o = xr.concat(dss, dim=pd.Index(ens, name='en'))

    # save to nc
    encoding = {v:{'_FillValue': None, 'dtype': 'float32'}
        for v in ['HU', 'TS', 'PDI', 'MDR', 'TROP']}
    encoding['year'] = {'_FillValue': None, 'dtype': 'int32'}
    encoding['en'] = {'_FillValue': None, 'dtype': 'int32'}
    ds_o.to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)
    
    t1 = datetime.datetime.now()
    print('[end]:', t1.strftime(tformat))
    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
    print()