import numpy as np
from maskedImage import MaskedImage
from nnf import Nnf

class Inpaint:
    def __init__(self,input,mask,radius):
        self.initial = MaskedImage(input,mask)
        self.radius = radius
        self.pyramid = [self.initial]
        source = self.initial
        while source.width>radius and source.height>radius:
            source = source.downsample()
            self.pyramid.append(source)
        maxLevel = len(self.pyramid)

        for level in range(maxLevel-1,0,-1):
            print(level)
            source = self.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
                self.nnfSourceToTarget = Nnf(source,target,radius)
                self.nnfSourceToTarget.randomize()

                self.nnfTargetToSource = Nnf(target,source,radius)
                self.nnfTargetToSource.randomize()
            else:
                newNnf = Nnf(source,target,radius)
                newNnf.initializeFromNnf(self.nnfSourceToTarget)
                self.nnfSourceToTarget = newNnf

                newNnfRev = Nnf(target,source,radius)
                newNnfRev.initializeFromNnf(self.nnfTargetToSource)
                self.nnfTargetToSource = newNnfRev
            target = self.ExpectationMaximization(level)
        return target.image

    def ExpectationMaximization(self,level):
        iterEM = 1+2*level
        iterNnf = min(7,1+level)
        source = self.nnfSourceToTarget.input
        target = self.nnfTargetToSource.output
        newTarget = None
        for emloop in range(1,iterEM+1):
            if (newTarget != None):
                self.nnfSourceToTarget.output = newTarget
                self.nnfTargetToSource.input = newTarget
                target = newTarget
                newTarget = None
            
            for y in range(source.height):
                for x in range(source.width):
                    if not source.containsMask(x,y,self.radius):
                        self.nnfSourceToTarget.field[y,x,0] = x
                        self.nnfSourceToTarget.field[y,x,1] = y
                        self.nnfSourceToTarget.field[y,x,2] = 0
            for y in range(target.height):
                for x in range(target.width):
                    if not source.containsMask(x,y,self.radius):
                        self.nnfTargetToSource.field[y,x,0] = x
                        self.nnfTargetToSource.field[y,x,1] = y
                        self.nnfTargetToSource.field[y,x,2] = 0
            self.nnfSourceToTarget.minimize(iterNnf)
            self.nnfTargetToSource.minimize(iterNnf)

            upscaled = False
            if level>=1 and emloop == iterEM:
                newSource = self.pyramid[level-1]
                newTarget = target.upscale(newSource.height,newSource.width)
                upscaled = True
            else:
                newSource = self.pyramid[level]
                newTarget = target.copy()
                upscaled = False

            vote = np.zeros((newTarget.width, newTarget.height, 4))
            self.ExpectationStep(self.nnfSourceToTarget,True,vote,newSource,upscaled)
            self.ExpectationStep(self.nnfTargetToSource,False,vote,newSource,upscaled)

            self.MaximizationStep(newTarget, vote)
            result = MaskedImage.resize(newTarget.image,self.initial.width,self.initial.height)
            from demo import Demo
            Demo.display(result)
        return newTarget

    def ExpectationStep(self,nnf,sourceToTarget, vote, source, upscale):
        for y in range(nnf.input.height):
            for x in range(nnf.input.width):
                xp = nnf.field[y,x,0]
                yp = nnf.field[y,x,1]
                dp = nnf.field[y,x,2]
                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:
                            self.weightedCopy(source,2*xs,2*ys,vote,2*xt,2*yt,w)
                            self.weightedCopy(source,2*xs+1,2*ys,vote,2*xt+1,2*yt,w)
                            self.weightedCopy(source,2*xs,2*ys+1,vote,2*xt,2*yt+1,w)
                            self.weightedCopy(source,2*xs+1,2*ys+1,vote,2*xt+1,2*yt+1,w)
                        else:
                            self.weightedCopy(source,xs,ys,vote,xt,yt,w)
    
    def weightedCopy(self,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(self,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,0] = r
                    target.image[y,x,1] = g
                    target.image[y,x,2] = b
                    target.mask[y,x] = False