import matplotlib.pyplot as plt
import numpy as np
import cv2

from maskedImage import MaskedImage
from nnf import Nnf

def read(file):
    img = plt.imread(file)
    if img.dtype == np.float32:
        img = (img * 255).astype(np.uint8)
        img = img[:,:,0:3]
    return img

def doTheInpainting(img,mask,radius):
    def maximizeForTheScale(scale):
        iterEM = 1+2*scale
        iterNnf = min(7,1+scale)
        source = sourceToTarget.input
        target = targetToSource.output
        newTarget = None
        for emloop in range(1,iterEM+1):
            if (newTarget != None):
                sourceToTarget.output = newTarget
                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):
                        sourceToTarget.field[y,x] = (x,y,0)
                        targetToSource.field[y,x] = (x,y,0)
            sourceToTarget.minimize(iterNnf)
            targetToSource.minimize(iterNnf)

            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

            vote = np.zeros((newTarget.width, newTarget.height, 4))
            ExpectationStep(sourceToTarget,True,vote,newSource,upscaled)
            ExpectationStep(targetToSource,False,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)
    radius = radius
    pyramid = [initial]
    source = initial
    while source.width>radius and source.height>radius:
        source = source.downsample()
        pyramid.append(source)
    maxLevel = len(pyramid)

    for level in range(maxLevel-1,0,-1):
        source = pyramid[level]
        if (level == maxLevel-1):
            target = source.copy()
            for y in range(target.height):
                for x in range(target.width):
                    target.mask[y,x] = False
            sourceToTarget = Nnf(source,target,radius)
            sourceToTarget.randomize()

            targetToSource = Nnf(target,source,radius)
            targetToSource.randomize()
        else:
            newNnf = Nnf(source,target,radius)
            newNnf.initializeFromNnf(sourceToTarget)
            sourceToTarget = newNnf

            newNnfRev = Nnf(target,source,radius)
            newNnfRev.initializeFromNnf(targetToSource)
            targetToSource = newNnfRev
        target, sourceToTarget, targetToSource = maximizeForTheScale(level)
        plt.imshow(target.image)
        plt.pause(0.01)
    return target.image

def ExpectationStep(nnf,sourceToTarget, vote, source, upscale):
    for y in range(nnf.input.height):
        for x in range(nnf.input.width):
            xp, yp, dp = nnf.field[y,x]
            w = MaskedImage.similarity[dp]
            for dy in range(-nnf.patchSize,nnf.patchSize):
                for dx in range(-nnf.patchSize,nnf.patchSize):
                    if sourceToTarget:
                        xs = x+dx
                        ys = y+dy
                        xt = xp+dx
                        yt = yp+dy
                    else:
                        xs = xp+dx
                        ys = yp+dy
                        xt = x+dx
                        yt = y+dy
                    if not 0<=xs<nnf.input.width:
                        continue
                    if not 0<=ys<nnf.input.height:
                        continue
                    if not 0<=xt<nnf.input.width:
                        continue
                    if not 0<=yt<nnf.input.height:
                        continue
                    if upscale:
                        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)

def weightedCopy(src,xs,ys,vote,xd,yd,w):
    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):
    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