from matplotlib.widgets import RectangleSelector
import matplotlib.pyplot as plt
import numpy as np

def doKnn(img,x1,y1,x2,y2):
    
    def getNotInBoundNeighbour(neighbour, x1,y1,x2,y2):
        mask = np.logical_or(
            np.logical_or(neighbour[:, 0] < y1, neighbour[:, 0] > y2),
            np.logical_or(neighbour[:, 1] < x1, neighbour[:, 1] > x2)
        )
        return neighbour[mask]

    def neighbourReelPixel(x,y):
        tNeighbour = np.copy(neighbour)
        tNeighbour = tNeighbour + np.array([y,x])
        return tNeighbour
    
    def getAvgPixelFromNeighbour(neighbour):
        return np.mean(img[neighbour[:,0],neighbour[:,1]], axis=0)
    
    neighbour = np.array([[-1,-1],[-1,0],[0,-1],[-1,1],[1,-1],[0,1],[1,0],[1,1]])
    x1c = x1
    y1c = y1
    x2c = x2
    y2c = y2
    # tant que les pixels en périphérie du trou ne se rejoignent pas alors le trou n'est pas comblé
    while x1 != x2 and y1 != y2:
        # on comble les pixels à gauche et à droite
        for x in range(x1,x2):
            currentNeighbour1 = neighbourReelPixel(x,y1)
            currentNeighbour2 = neighbourReelPixel(x,y2)
            currentNeighbour1 = getNotInBoundNeighbour(currentNeighbour1,x1,y1,x2,y2)
            currentNeighbour2 = getNotInBoundNeighbour(currentNeighbour2,x1,y1,x2,y2)
            currentColor1 = getAvgPixelFromNeighbour(currentNeighbour1)
            currentColor2 = getAvgPixelFromNeighbour(currentNeighbour2)
            img[y1,x] = currentColor1
            img[y2,x] = currentColor2
        # puis en haut et en bas
        for y in range(y1,y2):
            currentNeighbour1 = neighbourReelPixel(x1,y)
            currentNeighbour2 = neighbourReelPixel(x2,y)
            currentNeighbour1 = getNotInBoundNeighbour(currentNeighbour1,x1,y1,x2,y2)
            currentNeighbour2 = getNotInBoundNeighbour(currentNeighbour2,x1,y1,x2,y2)
            currentColor1 = getAvgPixelFromNeighbour(currentNeighbour1)
            currentColor2 = getAvgPixelFromNeighbour(currentNeighbour2)
            img[y,x1] = currentColor1
            img[y,x2] = currentColor2
        x1 += 1
        x2 -= 1
        y1 += 1
        y2 -= 1
    
    for x in range(x1c, x2c):
        for y in range(y1c, y2c):
            currentNeighbour = neighbourReelPixel(x, y)
            currentNeighbour = getNotInBoundNeighbour(currentNeighbour,0,0,0,0)
            currentColor = getAvgPixelFromNeighbour(currentNeighbour)
            img[y, x] = currentColor
    img[y1:y2,x1:x2] 
    return img

img = plt.imread('asset/boat.png')

if len(img.shape) == 2:
    img = np.stack((img,)*3, axis=-1)

def onselect(eclick, erelease):
    x1, y1 = eclick.xdata, eclick.ydata
    x2, y2 = erelease.xdata, erelease.ydata

    img_copy = np.copy(img)
    res = doKnn(img_copy,int(x1),int(y1),int(x2),int(y2))
    ax.imshow(res)
    plt.draw()

fig, ax = plt.subplots()
ax.imshow(img)
toggle_selector = RectangleSelector(ax, onselect,  useblit=True,
                                    button=[1], minspanx=5, minspany=5, spancoords='pixels',
                                    interactive=True)
plt.axis('off')
plt.show()