import xarray as xr
import numpy as np

# load data
wrf = xr.load_dataset("wrf_t2.nc")
weights = xr.load_dataset("wrf_to_era_weights.nc")

t2 = wrf["T2"].to_numpy()
new_t2 = np.zeros((19, 34), dtype=np.float32)

dst_add = weights["dst_address"].to_numpy()
src_add = weights["src_address"].to_numpy()
map_wts = weights["remap_matrix"].to_numpy().squeeze()
dst_mask = weights["dst_grid_imask"].to_numpy()
src_mask = weights["src_grid_imask"].to_numpy()
dst_frac = weights["dst_grid_frac"].to_numpy()

# bilinear interpolation
for n in range(int(weights.dims["num_links"] / 4)):
    noff = n * 4

    if (
        dst_mask[dst_add[noff]] and
        src_mask[src_add[noff]] and
        src_mask[src_add[noff + 1]] and
        src_mask[src_add[noff + 2]] and
        src_mask[src_add[noff + 3]] and
        dst_frac[dst_add[noff]] > 0
    ):
        new_t2[np.unravel_index(dst_add[noff], new_t2.shape)] = (
            t2[0][np.unravel_index(src_add[noff], t2[0].shape)] * map_wts[1 * noff] +
            t2[0][np.unravel_index(src_add[noff + 1], t2[0].shape)] * map_wts[1 * (noff + 1)] +
            t2[0][np.unravel_index(src_add[noff + 2], t2[0].shape)] * map_wts[1 * (noff + 2)] +
            t2[0][np.unravel_index(src_add[noff + 3], t2[0].shape)] * map_wts[1 * (noff + 3)]
        )

# change zeros to missin values
new_t2[new_t2 == 0] = np.nan

# create dataset
ds = xr.Dataset(data_vars={
    "T2": (
        ('time', 'latitude', 'longitude'),
        new_t2[None]
    )
}, coords={
    "latitude": np.arange(57.75, 53, -0.25),
    "longitude": np.arange(33.5, 42, 0.25),
})

# save netCDF
ds.to_netcdf("pyera2.nc")
