from __future__ import division

from numba import jit
import numpy as np

from imgProcessor.utils.growPositions import growPositions


def interpolate2dStructuredFastIDW(grid, mask, kernel=15, power=2,
                                   minnvals=5):
    '''
    FASTER IMPLEMENTATION OF interpolate2dStructuredIDW

    replace all values in [grid] indicated by [mask]
    with the inverse distance weighted interpolation of all values within
    px+-kernel
    [power] -> distance weighting factor: 1/distance**[power]

    [minvals] -> minimum number of neighbour values to find until
                 interpolation stops

    '''
    indices, dist = growPositions(kernel)
    weights = 1 / dist**(0.5 * power)

    return _calc(grid, mask, indices, weights, minnvals - 1)


@jit(nopython=True)
def _calc(grid, mask, indices, weights, minnvals):
    s0 = grid.shape[0]
    s1 = grid.shape[1]

    # FOR EVERY PIXEL
    for i in range(s0):
        for j in range(s1):

            if mask[i, j]:
                sumWi = 0.0
                value = 0.0
                c = 0
                # FOR EVERY NEIGHBOUR IN KERNEL
                for n in range(len(indices)):
                    iii = i + indices[n, 0]
                    jjj = j + indices[n, 1]
                    # if in image
                    if 0 <= iii < s0 and 0 <= jjj < s1:
                        if not mask[iii, jjj]:
                            wi = weights[n]
                            sumWi += wi
                            value += wi * grid[iii, jjj]
                            if c == minnvals:
                                break
                            c += 1
                    # too far out of the image:
                    elif (c > 0 and
                          (iii < -1 or iii > s0 + 1) and
                          (jjj < -1 or jjj > s1 + 1)):
                        break
                if sumWi:
                    grid[i, j] = value / sumWi

    return grid


if __name__ == '__main__':
    import pylab as plt
    import sys
    from time import time

    shape = (200, 200)

    # array with random values:
    arr = np.random.rand(*shape)
    # mask containing valid cells:
    mask = np.random.randint(50, size=shape)
    mask[mask < 1] = False
    mask = mask.astype(bool)

    # substituting all cells with mask==True with interpolated value:
    t0 = time()
    arr1 = interpolate2dStructuredFastIDW(arr.copy(), mask, kernel=20, power=1)
    arr2 = interpolate2dStructuredFastIDW(arr.copy(), mask, kernel=20, power=2)
    arr3 = interpolate2dStructuredFastIDW(arr.copy(), mask, kernel=20, power=3)
    arr5 = interpolate2dStructuredFastIDW(arr.copy(), mask, kernel=20, power=5)
    print('time=%f' % (time() - t0))

    if 'no_window' not in sys.argv:
        plt.figure('power=1')
        arr1[~mask] = np.nan
        plt.imshow(arr1)

        plt.figure('power=2')
        arr2[~mask] = np.nan

        plt.imshow(arr2)
        plt.figure('power=3')
        arr3[~mask] = np.nan
        plt.imshow(arr3)

        plt.figure('power=5')
        arr5[~mask] = np.nan
        plt.imshow(arr5)

        plt.show()
