import xarray as xr
from glob import glob
import os

## Create WRF GCM dataset for wind and temp bias-corrected by GridMET dataset
gcms = [
    #"access1.0_RCP45",
    #"access1.0_RCP85",
    #"access1.3_RCP85",
    #"bcc-csm1.1_RCP85",
    #"canesm2_RCP85",
    #"ccsm4_RCP85",
    #"csiro-mk3.6.0_RCP85",
    #"fgoals-g2_RCP85",
    #"gfdl-cm3_RCP85",
    #"giss-e2-h_RCP85",
    #"miroc5_RCP85",
    #"mri-cgcm3_RCP85",
    #"noresm1-m_RCP85",
    "pnnl"
]


idir = "/home/disk/rocinante/DATA/temp/WRF/var/"
odir = "/home/disk/rocinante/DATA/temp/crystal_fire/wrf/pnnl/"

# Helper function to open and format dataset
def open_ds(idir, gcm, var):
    print("\tOpening dataset: ", gcm, var)
    data = sorted(glob("{}/{}/2d/{}_year/*.nc".format(idir, gcm, var)))
    ds = xr.open_mfdataset(data)
    ds = ds.sel(x=slice(51,75), y=slice(53,91))
    #if var == "T2":
    #    ds = ds.resample(times='1D').max()
    #else:
    #    ds = ds.resample(times='1D').mean()
    print("\tDataset loaded...")
    return ds


print('Opening GridMET dataset')

for gcm in gcms:
    print("\nWorking on: ", gcm)
    gdir = "{}/{}/".format(odir, gcm)
    os.makedirs(gdir, exist_ok=True)

    # Open data, merge and save daily data
    tds = open_ds(idir, gcm, "T2")
    #uds = open_ds(idir, gcm, "U10")
    #vds = open_ds(idir, gcm, "V10")
    wds = open_ds(idir, gcm, "WND")
    tds['T2'] -= 273.15
    #tds['U10'] = uds.U10
    #tds['V10'] = vds.V10
    tds['WND'] = wds.WND
    
    ds = tds[['WND']].resample(times='1D').mean()
    ds['T2'] = tds.T2.resample(times='1D').max()
    
    print('\tSaving dataset')
    ds.to_netcdf("{}/{}.nc".format(gdir, gcm)) # save data

    # Clean-up
    tds.close()
    #uds.close()
    #vds.close()
    wds.close()