#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb 22 10:57:20 EST 2021
##SBATCH --job-name=array-job     # create a short name for your job
##SBATCH --output=slurm-%A_%a.out # STDOUT file
##SBATCH --error=slurm-%A_%a.err  # STDERR file
#SBATCH --nodes=1                # node count
#SBATCH --ntasks=1               # total number of tasks across all nodes
#SBATCH --cpus-per-task=20        # cpu-cores per task (>1 if multi-threaded tasks)
#SBATCH --mem-per-cpu=4G         # memory per cpu-core (4G is default)
#SBATCH --time=01:00:00          # total run time limit (HH:MM:SS)
##SBATCH --array=1-100              # job array with index values 1, 2, ..., 100
#SBATCH --mail-type=all          # send email on job start, end and fault
#SBATCH --mail-user=wenchang@princeton.edu
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
from itertools import product
from geopotential import geopotential_height_make
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
volcs = ['Pinatubo', 'Agung', 'StMaria', 'Chichon', 'Chichon_CMIP5volc'][3:4]
volcs = volcs
expnames = [f'{volc}_PI_ens_noleap' for volc in volcs] 
compname = 'atmos_month'
#danames = ['netrad_toa', 't_surf', 'precip']
danames = ['t_surf',]
#danames = danames[1:2]
ens = range(1, 31)
#ens = range(1, 2) #only select the first ensemble member (long run)
n_ens = len(ens)
nyears_extended = 10#20#10 is default
thisdir = os.path.dirname(__file__)
#thisdir = os.getcwd() #for sbatch
 
def get_volc_data(expname, compname, daname, ens, plev=None, extended=True):
    '''extract a variable from volcanic dataset and combine all the ensemble members'''
    print(daname, expname, compname)
    ofile = os.path.join(thisdir, f'{daname}.{expname}.{compname}.nens{n_ens}.nc')
    if nyears_extended != 10:
        ofile = ofile.replace('.nc', f'.nyears{5+nyears_extended}.nc')
    if plev:
        ofile = ofile.replace('.nc', f'.{plev}hPa.nc')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return

    dirin = f'/tigress/wenchang/MODEL_OUT/{expname}'
    dirin_extend = dirin + '_extend'

    # combine ensemble members
    das = []
    for en in ens:
        print(f'en = {en:02d} of {n_ens:02d}')
        #model output files and the extended
        #ifiles = glob.glob(f'{dirin}/en{en:02d}/POSTP/*.{compname}.nc')
        #ifiles.sort()
        #ifiles_extend = glob.glob(f'{dirin_extend}/en{en:02d}/POSTP/*.{compname}.nc')
        #ifiles_extend.sort()
        #print(ifiles + ifiles_extend)
        ifiles = f'{dirin}/en{en:02d}/POSTP/*.{compname}.nc'
        ds = xr.open_mfdataset(ifiles, combine='by_coords')
        if daname == 'z_full':
            da = geopotential_height_make(ds)
        else:
            da = ds[daname]
        if plev:
            #da = da.sel(pfull=plev, method='nearest')
            da = da.interp(pfull=plev) # WY: use interp to get the exact level
        print('loading...')
        da.load()
        if extended:
            ifiles = f'{dirin_extend}/en{en:02d}/POSTP/*.{compname}.nc'
            ds = xr.open_mfdataset(ifiles, combine='by_coords')
            if n_ens > 1:
                ds = ds.isel(time=slice(0, nyears_extended*12)) #only select the first 10 years of extended data (the first ensemble member has longer extended period)
            if daname == 'z_full':
                da_extend = geopotential_height_make(ds)
            else:
                da_extend = ds[daname]
            if plev:
                #da_extend = da_extend.sel(pfull=plev, method='nearest')
                da_extend = da_extend.interp(pfull=plev) # WY: use interp to get the exact level
            print('loading...')
            da_extend.load()
        da = xr.concat([da, da_extend], dim='time')#combine both raw and extended data over the time dimension
        das.append(da)
    print('concatenating...')
    dae = xr.concat(das, pd.Index(ens, name='en'))

    # convert to Python datetime
    new_time = [datetime.datetime(*t.timetuple()[0:6]) 
                            for t in dae.time.values]
    dae['time'] = new_time
    
    new_dims = dict(grid_xt='lon', grid_yt='lat')
    encoding = {daname: dict(zlib=True, complevel=1)}
    dae.rename(new_dims).to_dataset().to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)

for expname, daname in product(expnames, danames):
    get_volc_data(expname=expname, compname=compname, daname=daname, ens=ens)

if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()