import xarray as xr
import xesmf as xe
import sys


idir = "/home/disk/becassine/jswon11/DATA/WRF-daily/WRF-NARR/WRF-NARR_"
def read_WRF(var, styr, edyr):
    print(var, styr, edyr)
    ifile = f'{idir}{var}.zarr'
    print(ifile)
    ds = xr.open_dataset(ifile, engine='zarr')
    ds = ds.sel(time=slice(f'{styr}/01/01', f'{edyr}/12/31'))
    ds = ds.sel(x=slice(59,88),y=slice(153,190))
    ds = ds.resample(time='1M')

    if var == 'PREC':
        ds = ds.sum()
    else:
        ds = ds.mean()
        
    return ds
    
    
#if ((lat > 47) & (lat < 47.6) & (lon > -122.3) & (lon < -121.2)):

styr = 1991
edyr = 2020

ds_prec = read_WRF('PREC', styr, edyr)
ds_tavg = read_WRF('TAVG', styr, edyr)
ds_tmax = read_WRF('TMAX', styr, edyr)
ds_tmin = read_WRF('TMIN', styr, edyr)


odir = '/home/disk/rocinante/DATA/temp/kcp3/scripts/met_compare/green/narr/'
out = f'{odir}/WRF-NARR_{styr}-{edyr}.nc'
ds = xr.merge([ds_prec, ds_tavg, ds_tmax, ds_tmin])
print(out)
ds.to_netcdf(out)