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

# verrifier que l'image est de même taille que le mask

def onSelect(eclick, erelease):
    # remplace les pixels sélectioné en rouge et enregistre ses pixels comme étant à modifé
    x1, y1 = int(eclick.xdata), int(eclick.ydata)
    x2, y2 = int(erelease.xdata), int(erelease.ydata)
    img[y1:y2+1,x1:x2+1] = (255,0,0)
    mask[y1:y2+1,x1:x2+1] = True
    maskImg[y1:y2+1,x1:x2+1] = 255
    ax.imshow(img)
    plt.draw()

imgOr = read(input("lien relatif de l'image: "))
img = imgOr.copy()

rep = input("avez vous le mask de l'image (y:n): ")
yes = ["y","Y","yes","YES","Yes","o","O","oui","OUI"]
height, width = img.shape[:2]
if (rep in yes):
    # récupère le mask enregistré
    maskImg = read(input("lien relatife de l'image: "))
    if (len(maskImg.shape) == 3):
        maskImg = maskImg[:,:,0]
    mask = maskImg == 255
    img[mask] = (255,0,0)
else:
    # lance la fenêtre de selection
    fig, ax = plt.subplots()
    plt.axis("off")
    ax.imshow(img)
    mask = np.zeros((height, width), dtype=bool)
    maskImg = np.zeros((height, width), dtype=np.uint8)
    # permet la selection en rectangle sur l'image et appelle la fonction onSelect 
    toggle_selector = RectangleSelector(ax, onSelect,  useblit=True,
                                        button=[1], minspanx=5, minspany=5, spancoords='pixels',
                                        interactive=True)
    plt.imshow(img)
    plt.show()

# lance une nouvelle fenêtre que l'on vas pouvoir annimer en fonction de l'avancement de l'inpainting
plt.ion()
plt.axis("off")

plt.imshow(mask)
plt.show()
plt.title("mask de l'image")
plt.pause(1)
plt.title("génération de l'image")
# lance l'impainting sur img avec le mask mask avec une taile de patch de 2
res = doTheInpainting(img,mask,2)
plt.ioff()

if input("afficher le résultat avant après (y/n): ") in yes:
    fig, (ax1, ax2) = plt.subplots(1, 2)
    ax1.imshow(imgOr)
    ax1.set_title('Original')
    ax1.axis("off")
    ax2.imshow(res)
    ax2.set_title('Result')
    ax2.axis("off")
    plt.show()

if input("enregister le mask (y/n):") in yes:
    save_path = input("Entrez le chemin d'enregistrement du mask: ")
    plt.imsave(save_path, maskImg, cmap='gray')

if input("enregister l'image produite (y/n):") in yes:
    save_path = input("Entrez le chemin d'enregistrement de l'image: ")
    plt.imsave(save_path, res)