import numpy as np

def nn_interp(input_arr):

    shape = np.shape(input_arr)
    n_shape = np.size(shape)
    
    if(n_shape < 2):
        print('error. cant do nn_interp on arrays with rank < 2')
        sys.exit()
    
    else:

        iter_arr = np.ma.copy(input_arr)
        
        while( np.count_nonzero(iter_arr.mask) > 0 ):
            for shift in (-1,1):
                for axis in (-1,-2):        
                    arr_shifted=np.roll(iter_arr,shift=shift,axis=axis)
                    idx=~arr_shifted.mask * iter_arr.mask
                    iter_arr[idx]=arr_shifted[idx]

    return iter_arr



def nn_interp_1d(input_arr):

    shape = np.shape(input_arr)
    n_shape = np.size(shape)

    iter_arr = np.ma.copy(input_arr)
        
    while( np.count_nonzero(iter_arr.mask) > 0 ):
        for shift in (-1,1):
            arr_shifted=np.roll(iter_arr,shift=shift,axis=-1)
            idx=~arr_shifted.mask * iter_arr.mask
            iter_arr[idx]=arr_shifted[idx]

    return iter_arr