from imageRelation import ImageRelation
from maskedImage import MaskedImage
import matplotlib.pyplot as plt
from tkinter import filedialog
import concurrent.futures
import tkinter as tk
import numpy as np
import pyautogui
import cv2


def read(file):
    # fonction qui renvoit une image à partir de son chemin
    img = plt.imread(file)
    if img.dtype == np.float32: # si on est en valeur entre 0 et 1 au lieux de entre 0 et 255
        img = (img * 255).astype(np.uint8)
        img = img[:,:,0:3] # si on est en rgba
    return img

def doTheInpainting(img,mask,radius):
    # fonction qui effectue l'impainting
    def maximizeForTheScale(scale):
        # fonction interne qui effectue l'impainting pour une version réduite de l'image
        iterEM = 1+2*scale
        npPass = min(7,1+scale)
        source = sourceToTarget.input
        target = targetToSource.output
        newTarget = None
        for emloop in range(1,iterEM+1): # on fait les différentes pass pour augmenter petit à petit la qualité de l'image
            # initialisation
            if (newTarget != None):
                targetToSource.input = newTarget
                target = newTarget
                newTarget = None
            for y in range(source.height):
                for x in range(source.width):
                    if not source.containsMask(x,y,radius):
                        targetToSource.field[y,x] = (x,y,0)
            # on cherche les patchs
            targetToSource.findBestPatch(npPass)
            # on crée la source et la target
            upscaled = False
            if scale>=1 and emloop == iterEM:
                newSource = images[scale-1]
                newTarget = target.upScale(newSource.height,newSource.width)
                upscaled = True
            else:
                newSource = images[scale]
                newTarget = target.copy()
                upscaled = False
            # on vote, on applique les votes puis on affiche les changements
            vote = np.zeros((newTarget.width, newTarget.height, 4))
            voteForPixels(targetToSource,vote,newSource,upscaled)
            applyVote(newTarget, vote)
            result = cv2.resize(newTarget.image, (initial.width, initial.height), interpolation=cv2.INTER_AREA)
            plt.imshow(result)
            plt.pause(0.01)
        return newTarget, sourceToTarget, targetToSource

    initial = MaskedImage(img,mask)
    images = [initial]
    source = initial
    while source.width>radius and source.height>radius: # crée les versions réduite en qualité de l'image
        source = source.downScale()
        images.append(source)
    maxLevel = len(images)

    for level in range(maxLevel-1,0,-1): # pour toutes les tailes de l'image, progressivement completer le trou à partir des images de qualité infèrieur déjà calculé
        source = images[level]
        if (level == maxLevel-1):
            target = source.copy()
            target.mask[0:target.height,0:target.width] = False
            sourceToTarget = ImageRelation(source,target,radius)
            sourceToTarget.randomize()

            targetToSource = ImageRelation(target,source,radius)
            targetToSource.randomize()
        else:
            newImRel = ImageRelation(source,target,radius)
            newImRel.initializeFromImageRelation(sourceToTarget)
            sourceToTarget = newImRel

            newImRelRev = ImageRelation(target,source,radius)
            newImRelRev.initializeFromImageRelation(targetToSource)
            targetToSource = newImRelRev
        target, sourceToTarget, targetToSource = maximizeForTheScale(level)
        plt.imshow(target.image)
        plt.pause(0.01)
    return target.image

def voteForPixels(imRel, vote, source, upscale):
    def voteForNb(nb):
        # ajoute au vote pour toutes les pixels les valeurs déterminé lors de la recherche des patchs
        wid = imRel.input.width//7
        for y in range(imRel.input.height):
            for x in range(nb*wid,(nb+1)*wid if nb != 7 else imRel.input.width): # divise par 8 l'image dans la largeur pour calculer ses 8 parties en même temps grâce aux threads
                xp, yp, dp = imRel.field[y,x]
                w = MaskedImage.similarity[dp]
                for dy in range(-imRel.patchSize,imRel.patchSize): # pour toutes les pixels du patch qui ne sorte pas de l'image
                    ys = yp+dy
                    if not 0<=ys<imRel.input.height:
                        continue
                    yt = y+dy
                    if not 0<=yt<imRel.input.height:
                        continue
                    for dx in range(-imRel.patchSize,imRel.patchSize):
                        xs = xp+dx
                        if not 0<=xs<imRel.input.width:
                            continue
                        xt = x+dx
                        if not 0<=xt<imRel.input.width:
                            continue
                        if upscale: # si on change d'échelle prendre les 4 pixels autour pour faire comme un knn
                            addToVote(source,2*xs,2*ys,vote,2*xt,2*yt,w)
                            addToVote(source,2*xs+1,2*ys,vote,2*xt+1,2*yt,w)
                            addToVote(source,2*xs,2*ys+1,vote,2*xt,2*yt+1,w)
                            addToVote(source,2*xs+1,2*ys+1,vote,2*xt+1,2*yt+1,w)
                        else:
                            addToVote(source,xs,ys,vote,xt,yt,w)
    # utilise une pool pour profiter des 8 threads de ma machine pour réduire le temp d'execution
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=8)
    for i in range(8):
        pool.submit(voteForNb,i)
    pool.shutdown(wait=True)

def addToVote(src,xs,ys,vote,xd,yd,w):
    # on ajoute au vote d'une pixel la couleur de la pixel asigné multiplié par un poid qui change en fonction de la similarité entre les deux pixels
    if src.mask[ys,xs]:
        return
    vote[xd,yd,0] += w*src.image[ys,xs,0]
    vote[xd,yd,1] += w*src.image[ys,xs,1]
    vote[xd,yd,2] += w*src.image[ys,xs,2]
    vote[xd,yd,3] += w

def applyVote(target,vote):
    # appliquer les changements de pixels déterminé par les votes
    for y in range(target.height):
        for x in range(target.width):
            if vote[x,y,3]>0:
                r = int(vote[x,y,0]/vote[x,y,3])
                g = int(vote[x,y,1]/vote[x,y,3])
                b = int(vote[x,y,2]/vote[x,y,3])
                target.image[y,x] = (r,g,b)
                target.mask[y,x] = False


def binDialog(titre,text):
    # box de dialogue à choix bianire
    pyautogui.FAILSAFE = True
    return pyautogui.confirm(
        text=text,
        title=titre,
        buttons=['Oui', 'Non']
    )

def dialog(titre,text):
    # box de dialogue à choix bianire
    pyautogui.FAILSAFE = True
    pyautogui.confirm(
        text=text,
        title=titre,
        buttons=['Ok']
    )

def stringDialog(titre, text):
    # box de dialogue avec un input
    pyautogui.FAILSAFE = True
    return pyautogui.prompt(
        text=text,
        title=titre,
        default=''
    )

def selectImage():
    # fonction de selection d'une image
    root = tk.Tk()
    root.withdraw()  # Hide the main window
    file_path = filedialog.askopenfilename(
        title="Select an image",
        filetypes=[
            ("Image files", "*.png *.jpg *.jpeg *.bmp")
        ]
    )
    return file_path