#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Nov 12 12:54:39 EST 2019
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
print()

data_name = 'ts'
ssp = 'ssp245'
#ifile_historical = f'{data_name}.historical.1850-2014.ens23.nc'
ifile_historical = f'{data_name}.historical.1850-2014.ens27.nc'
#ifile_ssp = f'{data_name}.{ssp}.2015-2100.ens16.nc'
ifile_ssp = f'{data_name}.{ssp}.2015-2100.ens19.nc'
 
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.datetime.now()
    print('[start]:', t0.strftime(tformat))

    # 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')
    
    t1 = datetime.datetime.now()
    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
    print('[end]:', t1.strftime(tformat))
    print()