#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Sep 18 21:10:30 EDT 2019
if __name__ == '__main__':
    from misc.timer import Timer
    _t = Timer(f'start {__file__}')
from datetime import datetime
import os.path, sys, os, glob
#import matplotlib.pyplot as plt
import xarray as xr, numpy as np, pandas as pd
#import geoxarray
from xlib.io import get_data
print()

expnames = [
    'Pinatubo_PI_ens_noleap',
    'Agung_PI_ens_noleap',
    'StMaria_PI_ens_noleap',
    'Chichon_PI_ens_noleap',
    'CTL1860_noleap_tigercpu_intelmpi_18_576PE',
    ]
data_names = ['swdn_sfc', 'swup_sfc', 'lwdn_sfc', 'lwup_sfc' ]
odata_name = 'netdn_sfc'
long_name = 'net downward radiation at surface'
do_clearsky = True
if do_clearsky:
    data_names = [dname+'_clr' for dname in data_names]
    odata_name = odata_name+'_clr'
    lon_name = 'clearsky ' + long_name


if __name__ == '__main__': 
    for expname in expnames: 
        ofile = f'{expname}.{odata_name}.nc'
        if expname.startswith('CTL1860'):
            ofile = f'{expname}.{odata_name}.ens.nc'
        if os.path.exists(ofile):
            print('[exists]:', ofile)
            continue
        das = dict()
        for data_name in data_names:
            ifile = f'{expname}.{data_name}.nc'
            if expname.startswith('CTL1860'):
                ifile = f'{expname}.{data_name}.ens.nc'
            das[data_name] = xr.open_dataarray(ifile)
        if do_clearsky:
            da = das['swdn_sfc_clr'] - das['swup_sfc_clr'] + das['lwdn_sfc_clr'] - das['lwup_sfc_clr']
        else:
            da = das['swdn_sfc'] - das['swup_sfc'] + das['lwdn_sfc'] - das['lwup_sfc']
        da.attrs['units'] = das[data_names[0]].attrs['units']
        da.attrs['long_name'] = long_name

        # save to netcdf
        ds = da.to_dataset(name=odata_name)
        encoding = {odata_name: {'dtype': 'float32', 'zlib': True, 'complevel': 1}}
        print('saving to netcdf ...')
        ds.load().to_netcdf(ofile, encoding=encoding)
        print('[saved]:', ofile)

    # the end
    _t.check(f'end {__file__}')