#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Apr 22 11:58:41 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
dsname = 'Tas'
daname = 'Tas'
zero_urban = False
if '0urban' in sys.argv: zero_urban = True

ifile = f'/tigress/wenchang/analysis/WWA/2021PNWheatwave/data_{dsname}_AM4urban_wasteCool_amip_5ens_1870-2020.nc'
if zero_urban: ifile = ifile.replace('wasteCool', 'wasteCool_0urban')
print(f'{ifile = }')
ibasename = os.path.basename(ifile)
ofile = ibasename.replace('5ens', '10ens')
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
ds = xr.open_dataset(ifile)

ifile = f'/tigress/wenchang/analysis/WWA/2021PNWheatwave/wytemp/data_{dsname}_AM4urban_wasteCool_amip_new5ens_1870-2020.nc'
if zero_urban: ifile = ifile.replace('wasteCool', 'wasteCool_0urban')
print(f'{ifile = }')
ds1 = xr.open_dataset(ifile)

ds = xr.concat([ds.sel(en=[0,21,22,23,24]), ds1], dim='en')
ds.to_netcdf(ofile, encoding={daname: {'zlib': True, 'complevel': 1}})
print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    plt.show()