import xarray as xr
import numpy as np

sites = {
    ('Blnghm',48.79389,-122.53722),
    ('Bremerton',47.48333,-122.76667),
    ('FridayHarbor',48.52222,-123.02306),
    ('OlyArpt',46.9733,-122.9033),
    ('OrcasArpt',48.70833,-122.91056),
    ('PortAngeles',48.141481,-123.414075),
    ('Quillayute',47.9375,-124.555),
    ('Quinault',47.5139,-123.812),
    ('SeaTac',47.4444,-122.3138),}


def get_latlon2d(ds, lat, lon):
    lats = ds.lat.values
    lons = ds.lon.values
    dist = ((lats - lat)**2 + (lons-lon)**2)**0.5
    i,j = np.unravel_index(np.argmin(dist, axis=None), dist.shape)
    print("[{}, {}]: {:0.5f}, {:0.5f}".format(i, j, lats[i,j], lons[i,j]))
    return (i, j)



def process_site(site, pds):
    #odir = '/home/disk/rocinante/DATA/temp/TNC_stormwater/ASOS/pnnl/'
    #odir = '/home/disk/rocinante/DATA/temp/TNC_stormwater/ASOS/bias-correction/'
    #odir = '/home/disk/rocinante/DATA/temp/TNC_stormwater/ASOS/prism/'
    odir = '/home/disk/rocinante/DATA/temp/TNC_stormwater/ASOS/bc_hourly/'
    
    (name, lat, lon) = site
    print(name, lat, lon)
    (x, y) = get_latlon2d(pds, lat, lon)
    cds = pds.sel(x=x, y=y)
    #out = '{}/{}_WRF-NARR.nc'.format(odir, name)
    out = '{}/{}_bc.nc'.format(odir, name)
    #out = '{}/{}_prism.nc'.format(odir, name)
    cds.to_netcdf(out)
    
    #df.to_csv('data/{}_ds3505.csv'.format(site), index=False, float_format='%0.2f')


if __name__ == '__main__':
    #pds = xr.open_dataset('/home/disk/rocinante/DATA/temp/TNC_stormwater/pnnl/pnnl_1981-2020.nc')    
    #pds = xr.open_dataset('/home/disk/rocinante/DATA/temp/TNC_stormwater/bias_correction/data/prism_interp.nc')
    #pds = xr.open_dataset('/home/disk/rocinante/DATA/temp/TNC_stormwater/bias_correction/data/bc_1.0bin.nc')
    #pds = xr.open_dataset('/home/disk/rocinante/DATA/temp/TNC_stormwater/bias_correction/data/bc_1.0bin.nc')
    
    for site in sites:
        process_site(site, pds)