import random
import numpy as np
from maskedImage import MaskedImage

class Nnf:
    def __init__(self,input,output,patchSize):
        self.input = input
        self.output = output
        self.patchSize = patchSize

    def randomize(self):
        self.field = np.zeros((self.input.height, self.input.width, 3), dtype=int)
        self.field[:,:,2] = MaskedImage.DSCALE
        self.field[:,:,0] = np.random.randint(0,self.output.width,(self.input.height,self.input.width))
        self.field[:,:,1] = np.random.randint(0,self.output.height,(self.input.height,self.input.width))
        self.initialize()
    
    def initializeFromNnf(self,nnf):
        self.field = np.zeros((self.input.height, self.input.width, 3), dtype=int)
        fx = int(self.input.width/nnf.input.width)
        fy = int(self.input.height/nnf.input.height)
        for y in range(self.input.height):
            for x in range(self.input.width):
                xl = min(int(x/fx),nnf.input.width-1)
                yl = min(int(y/fy),nnf.input.height-1)
                self.field[y,x] = (nnf.field[yl,xl,0]*fx, nnf.field[yl,xl,1]*fy, MaskedImage.DSCALE)
        self.initialize()
    
    def initialize(self):
        for y in range(self.input.height):
            for x in range(self.input.width):
                self.field[y,x,2] = self.distance(x,y,self.field[y,x,0],self.field[y,x,1])
                iter= 0
                maxIter = 10
                while (self.field[y,x,2] == MaskedImage.DSCALE and iter<maxIter):
                    self.field[y,x] = (random.randint(0,self.output.width),random.randint(0,self.output.height),self.distance(x,y,self.field[y,x,0],self.field[y,x,1]))
                    iter += 1
    
    def minimize(self,nbPass):  
        for i in range(nbPass):
            for y in range(self.input.height-1):
                for x in range(self.input.width-1):
                    if (self.field[y,x,2]>0):
                        self.minimizeLink(x,y,1)
            for y in range(self.input.height-1,0,-1):
                for x in range(self.input.width-1,0,-1):
                    if (self.field[y,x,2]>0):
                        self.minimizeLink(x,y,-1)
    
    def minimizeLink(self,x,y,direction):
        # horizontale
        if (0<x-direction<self.input.width):
            xp = self.field[y,x-direction,0] +direction
            yp = self.field[y,x-direction,1]
            dp = self.distance(x,y,xp,yp)
            if (dp<self.field[y,x,2]):
                self.field[y,x] = (xp,yp,dp)
        
        # verticale
        if (0<y-direction<self.input.height):
            xp = self.field[y-direction,x,0]
            yp = self.field[y-direction,x,1] + direction
            dp = self.distance(x,y,xp,yp)
            if (dp<self.field[y,x,2]):
                self.field[y,x] = (xp,yp,dp)
        
        # recherche random
        zoneRecherche = max(self.output.height,self.output.width)
        while zoneRecherche>0:
            xp = self.field[y,x,0] + random.randint(0,2*zoneRecherche)-zoneRecherche
            yp = self.field[y,x,1] + random.randint(0,2*zoneRecherche)-zoneRecherche
            xp = max(0,min(self.output.width-1,xp))
            yp = max(0,min(self.output.height-1,yp))
            dp = self.distance(x,y,xp,yp)
            if (dp<self.field[y,x,2]):
                self.field[y,x] = (xp,yp,dp)
            zoneRecherche = int(zoneRecherche/2)

    def distance(self,x,y,xp,yp):
        return distance(self.input,x,y,self.output,xp,yp,self.patchSize)

def distance(source, xs, ys, target, xt, yt, patchSize):
        ssd_max = 255 ** 2
        distance, wsum = 0, ssd_max * (patchSize*2)**2
        for dy in range(-patchSize, patchSize):
            yks, ykt = ys + dy, yt + dy
            if not (1 <= yks < source.height - 1 and 1 <= ykt < target.height - 1):
                distance += ssd_max * patchSize * 2
                continue
            for dx in range(-patchSize, patchSize):
                xks, xkt = xs + dx, xt + dx
                if not (1 <= xks < source.width - 1 and 1 <= xkt < target.width - 1):
                    distance += ssd_max
                    continue
                if source.containsMask(xks, yks,patchSize):
                    distance += ssd_max
                    continue
                if target.containsMask(xkt, ykt,patchSize):
                    distance += ssd_max
                    continue
                ssd = np.sum((source.image[yks, xks] - target.image[ykt, xkt]) ** 2)
                distance += ssd
        return int(MaskedImage.DSCALE * distance / wsum)