#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Nov 12 12:54:39 EST 2019
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, datetime, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
if __name__ == '__main__':
    tt.check('end import')
print()

data_name = 'ts'
ssp = 'ssp585'
#ifile_historical = f'{data_name}.historical.1850-2014.ens23.nc'
#ifile_historical = f'{data_name}.historical.1850-2014.ens27.nc'
ifiles = glob.glob(f'/tigress/wenchang/data/cmip6/variables/{data_name}/historical/wy_ens/{data_name}.historical.1850-2014.ens*.nc')
ifiles.sort()
ifile_historical = ifiles[-1]
if not os.path.exists( os.path.basename(ifile_historical) ):
    os.symlink(ifile_historical, os.path.basename(ifile_historical))
    print('**updated**:', ifile_historical)
#ifile_ssp = f'{data_name}.{ssp}.2015-2100.ens16.nc'
#ifile_ssp = f'{data_name}.{ssp}.2015-2100.ens19.nc'
ifiles = glob.glob(f'/tigress/wenchang/data/cmip6/variables/{data_name}/ssp585/wy_ens/{data_name}.{ssp}.2015-2100.ens*.nc')
ifiles.sort()
ifile_ssp = ifiles[-1]
if not os.path.exists( os.path.basename(ifile_ssp) ):
    os.symlink(ifile_ssp, os.path.basename(ifile_ssp))
    print('**updated**:', ifile_ssp)
 
if __name__ == '__main__':
    # open datasets
    print(ifile_historical,'...')
    da_historical = xr.open_dataset(ifile_historical)[data_name]
    print(ifile_ssp, '...')
    da_ssp = xr.open_dataset(ifile_ssp)[data_name]
    # models in common
    models = [m for m in da_ssp.model.values if m in da_historical.model.values]
    print(len(models), 'models:', models)
    # concat experiments
    ofile = f'{data_name}.historical_{ssp}.1850-2100.ens{len(models):02d}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        print()
        sys.exit()
    print('concat...')
    da = xr.concat([da_historical.sel(model=models), da_ssp.sel(model=models)], dim='time')
    # load data into memory
    print('loading...')
    ds = da.to_dataset(name=data_name).load()
    # save dataset
    print('saving...')
    encoding = {data_name: {'dtype': 'float32', 'zlib': True, 'complevel': 1}}
    ds.to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile, f'{os.path.getsize(ofile)/1e9:.1f}G')
    
    tt.check('**done**')
    print()