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


def patchMatch(img,x1,y1,x2,y2,patchSize=20, iterations=1000):
    height, width, _ = img.shape
    xx1 = max(0, x1-patchSize)
    yy1 = max(0, y1-patchSize)
    xx2 = min(width, x2+patchSize)
    yy2 = min(height, y2+patchSize)

    mask = np.ones((height, width), dtype=bool)
    mask[yy1:yy2, xx1:xx2] = False
    mask[y1:y2, x1:x2] = True
    img[int(y1):int(y2), int(x1):int(x2)] = np.mean(img[~mask], axis=(0, 1))
    img_copy = np.copy(img)

    div = 10

    if (xx2-xx1 < patchSize or yy2-yy1 < patchSize):
        return img

    def getRandomPatch():
        rx = np.random.randint(0, width - patchSize)
        ry = np.random.randint(0, height - patchSize)
        return rx, ry
    
    def fusion(patch1, patch2):
        return (patch1*2+patch2) / 3
    
    def distance(patch1, patch2):
        return np.sum((patch1 - patch2) ** 2)
    
    def gradientDescent(x, y, bestPatch, bestDistance):
        neighbors = [(-div,0), (div,0), (0,-div), (0,div), (-div,-div), (-div,div), (div,-div), (div,div)]
        patch = img[y:y + patchSize, x:x + patchSize]
        hasChanged = True
        while hasChanged:
            hasChanged = False
            for nx, ny in neighbors:
                cx = bestPatch[0] + nx
                cy = bestPatch[1] + ny
                if cx < 0 or cy < 0 or cx >= width - patchSize or cy >= height - patchSize:
                    continue

                neighborPatch = img[cy:cy + patchSize, cx:cx + patchSize]
                neighborDist = distance(patch, neighborPatch)
                if neighborDist < bestDistance:
                    hasChanged = True
                    bestPatch = [cx, cy]
                    bestDistance = neighborDist
        return bestPatch, bestDistance


    for x in range(xx1,xx2-patchSize,int(patchSize/div)):
        for y in range(yy1,yy2-patchSize,int(patchSize/div)):
            px, py = getRandomPatch()
            bestPatch = [px, py]
            bestDistance = distance(img[y:y+patchSize,x:x+patchSize], img[py:py+patchSize,px:px+patchSize])
            firstDistance = np.copy(bestDistance)
            for _ in range(iterations):
                px, py = getRandomPatch()
                currentDistance = distance(img[y:y+patchSize,x:x+patchSize], img[py:py+patchSize,px:px+patchSize])
                if currentDistance > firstDistance:
                    continue
                currentFirstDistance = np.copy(currentDistance)
                currentPatch, bestDistance = gradientDescent(x, y, bestPatch, bestDistance)
                if currentDistance < bestDistance:
                    firstDistance = currentFirstDistance
                    bestPatch = currentPatch
                    bestDistance = currentDistance
            img_copy[y:y+patchSize, x:x+patchSize] = fusion(img_copy[y:y+patchSize, x:x+patchSize],img[bestPatch[1]:bestPatch[1]+patchSize, bestPatch[0]:bestPatch[0]+patchSize])
            img[y1:y2,x1:x2] = img_copy[y1:y2,x1:x2]
    return img


# Load the image using matplotlib
img = plt.imread('simson.png')

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

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

    img_copy = np.copy(img)
    res = patchMatch(img_copy,int(x1),int(y1),int(x2),int(y2))
    ax.imshow(res)
    plt.draw()
    print("drawed")

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()