#!/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
expname = 'CTL1860_noleap_tigercpu_intelmpi_18_576PE'
#compname = 'atmos_month'
#danames = ['netrad_toa', 't_surf', 'precip']
compname = 'ocean'
danames = ['temp',]
year_start = 11
n_ens = 30
n_years = 15
#subset
lons = slice(-170,-120)
lats = slice(-40,40)
depths = slice(0, 300)
subset = f'{-lons.start}W-{-lons.stop}W_{-lats.start}S-{lats.stop}N_{depths.start}-{depths.stop}m'
func_subset = lambda x: x.sel(xt_ocean=lons, yt_ocean=lats, st_ocean=depths)

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'{daname}.{expname}.{compname}.nens{n_ens}.nyears{n_years}.{subset}.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}'
    
    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]
        #subset
        da = da.pipe(func_subset)
        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)
    dae.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()