#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb 22 10:07:16 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
from geopotential import geopotential_height_make
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#model = 'FLOR'
#expname = 'CTL1860_noleap_tigercpu_intelmpi_18_576PE'
model = 'AM2.5'
expname = 'CTL1860_florCTL1860noleapSST_tigercpu_intelmpi_18_540PE'
compname = 'atmos_month'
danames = ['netrad_toa', ]
year_start = 111
n_ens = 10
n_years = 5
thisdir = os.path.dirname(__file__)
def get_ctl_data(expname, compname, daname, year_start, n_ens, n_years=5, plev=None): 
    '''extract a varialbe from ctl experiment and tranform into ensembles'''
    ofile = os.path.join(thisdir, f'{model}.{daname}.{expname}.{compname}.nens{n_ens}.nyears{n_years}.nc')
    if plev:
        ofile = ofile.replace('.nc', f'.{plev}hPa.nc')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return

    if model in ('FLOR',):
        dirin = f'/tigress/wenchang/MODEL_OUT/{expname}'
    else:
        dirin = f'/tigress/wenchang/MODEL_OUT/{model}/{expname}'
    
    ens = range(1, n_ens+1)
    # combine ensemble members
    das = []
    for en in ens:
        print(f'en = {en:02d} of {n_ens:02d}')
        ifiles = [f'{dirin}/POSTP/{year:04d}0101.{compname}.nc'
            for year in range(year_start + en -1, year_start + en -1 + n_years)]
        #print(ifiles)
        ds = xr.open_mfdataset(ifiles)
        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()
        #reset the time coordinate the same across ensembles
        time_new = xr.cftime_range(start='2001-01-01', freq='MS', periods=n_years*12)#, calendar='noleap')
        da = da.assign_coords(time=time_new)
        das.append(da)
    print('concatanating...')
    dae = xr.concat(das, pd.Index(ens, name='en'))
    
    new_dims = dict(grid_xt='lon', grid_yt='lat')
    encoding = {daname: dict(zlib=True, complevel=1)}
    print('saving...')
    dae.rename(new_dims).to_dataset().to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)


for daname in danames:
    get_ctl_data(expname=expname, compname=compname, daname=daname, 
        year_start=year_start, n_ens=n_ens, n_years=n_years)

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()