from maskedImage import MaskedImage
import matplotlib.pyplot as plt
import concurrent.futures
from imageRelation import ImageRelation
import numpy as np
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 = pyramid[scale-1]
                newTarget = target.upscale(newSource.height,newSource.width)
                upscaled = True
            else:
                newSource = pyramid[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))
            ExpectationStep(targetToSource,vote,newSource,upscaled)
            MaximizationStep(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)
    pyramid = [initial]
    source = initial
    while source.width>radius and source.height>radius: # crée les versions réduite en qualité de l'image
        source = source.downsample()
        pyramid.append(source)
    maxLevel = len(pyramid)

    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 = pyramid[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 ExpectationStep(imRel, vote, source, upscale):
    def ExpectationStepForNb(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
                            weightedCopy(source,2*xs,2*ys,vote,2*xt,2*yt,w)
                            weightedCopy(source,2*xs+1,2*ys,vote,2*xt+1,2*yt,w)
                            weightedCopy(source,2*xs,2*ys+1,vote,2*xt,2*yt+1,w)
                            weightedCopy(source,2*xs+1,2*ys+1,vote,2*xt+1,2*yt+1,w)
                        else:
                            weightedCopy(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(ExpectationStepForNb,i)
    pool.shutdown(wait=True)

def weightedCopy(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 MaximizationStep(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