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




def patch_match(img, patch_size=3, iterations=1):
    height, width, _ = img.shape
    offsets = np.zeros((height, width, 2), dtype=np.int32)

    def random_offsets():
        return np.random.randint(-patch_size, patch_size + 1, size=(height, width, 2))

    def distance(patch1, patch2):
        return np.sum((patch1 - patch2) ** 2)

    def get_patch(x, y):
        return img[max(0, y):min(height, max(0, y) + patch_size), max(0, x):min(width, max(0, x) + patch_size)]

    offsets = random_offsets()

    for _ in range(iterations):
        for y in range(patch_size,height-patch_size*2):
            for x in range(patch_size,width-patch_size*2):
                best_offset = offsets[y, x]
                best_distance = distance(get_patch(x, y), get_patch(x + best_offset[0], y + best_offset[1]))
                for dy in range(-1, 2):
                    for dx in range(-1, 2):
                        if dx == 0 and dy == 0:
                            continue
                        new_offset = best_offset + [dx, dy]
                        if (new_offset[0]+x>width-patch_size or new_offset[1]+y>height-patch_size):
                            continue
                        new_distance = distance(get_patch(x, y), get_patch(x + new_offset[0], y + new_offset[1]))
                        if new_distance < best_distance:
                            best_distance = new_distance
                            best_offset = new_offset
                offsets[y, x] = best_offset
    
    result = np.zeros_like(img)
    for y in range(height):
        for x in range(width):
            offset = offsets[y, x]
            result[y, x] = img[(y + offset[1]) % height, (x + offset[0]) % width]

    return result





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


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

    avg_color = np.mean(img, axis=(0, 1))
    img_copy = np.copy(img)
    img_copy[int(y1):int(y2), int(x1):int(x2)] = avg_color
    res = patch_match(img_copy)
    ax.imshow(res)
    plt.draw()
    print("drawed")
    #ax.imshow(img_copy)
    #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()