import numpy as np
import cv2
import sys

# Used for the UI in drawing the uncertain and foreground areas
class RectangleImage:
    def __init__(self, img, progChan, compChan):
        self.origImg = img
        self.img = img.copy()
        self.updateImg = img.copy()
        self.progChan = progChan
        self.compChan = compChan
        self.rects = []
        self.start = None

    # Begin drawing current rectangle, should correspond to mouse down
    def startTemp(self, x, y):
        self.start = (y, x)

    # Continue drawing current rectangle, should correspond to mouse move
    def continueTemp(self, x, y):
        if self.start != None:
            self.updateImg = self.img.copy()
            topL = (min(self.start[0], y), min(self.start[1], x))
            botR = (max(self.start[0], y), max(self.start[1], x))
            self.updateImg[topL[0]:botR[0]+1,topL[1]:botR[1]+1,self.progChan] = 255

    # Continue drawing current rectangle, should correspond to mouse up
    def fixTemp(self, x, y):
        if self.start != None:
            topL = (min(self.start[0], y), min(self.start[1], x))
            botR = (max(self.start[0], y), max(self.start[1], x))
            self.img[topL[0]:botR[0]+1,topL[1]:botR[1]+1,self.compChan] = 255
            self.updateImg = self.img.copy()
            self.rects.append((topL, botR))
            self.start = None

    # Remove most recently draw rectangle and reset current, should correspond to undo
    def removeRect(self):
        if len(self.rects) > 0:
            self.start = None
            self.rects.pop()
            self.img = self.origImg.copy()
            for rect in self.rects:
                self.img[rect[0][0]:rect[1][0]+1,rect[0][1]:rect[1][1]+1,self.compChan] = 255
            self.updateImg = self.img.copy()

    # Event redirector
    def drawRect(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.startTemp(x, y)
        elif event == cv2.EVENT_MOUSEMOVE:
            self.continueTemp(x, y)
        elif event == cv2.EVENT_LBUTTONUP:
            self.fixTemp(x,y)
        cv2.imshow('Image', self.updateImg)

# Arguments
imageFile, bgImageFile = sys.argv[1:3]
# imageFile = 'image1.jpg'
# bgImageFile = 'image1bg.jpg'
img = cv2.imread(imageFile, 1)
bgImg = cv2.imread(bgImageFile, 1)
k = (int(sys.argv[3]) if len(sys.argv) >= 4 else 5)

# Define unknown rectangles
drawUncertain = RectangleImage(img, 1, 0)
cv2.namedWindow('Image')
cv2.imshow('Image', img)
cv2.setMouseCallback('Image', drawUncertain.drawRect)
loop = True
while(loop):
    key = cv2.waitKey(20)
    if key == 115: # s -> move on
        loop = False
    elif key == 122:  # z -> undo
        drawUncertain.removeRect()
unknownRects = drawUncertain.rects

# Define foreground rectangles
drawForegound = RectangleImage(drawUncertain.img, 2, 1)
cv2.setMouseCallback('Image', drawForegound.drawRect)
loop = True
while(loop):
    key = cv2.waitKey(20)
    if key == 115: # s -> move on
        loop = False
    elif key == 122:  # z -> undo
        drawForegound.removeRect()
    elif key == 97:
        print drawUncertain.rects
cv2.destroyAllWindows()
cv2.imwrite('output/ImageMarked.jpg', drawForegound.img)
foregroundRects = drawForegound.rects

# Convert from rectangle coordinates to assignment entries
assignments = np.zeros(img.shape[:-1], np.uint8)
for rect in unknownRects:
    assignments[rect[0][0]:rect[1][0]+1,rect[0][1]:rect[1][1]+1] = 100
for rect in foregroundRects:
    assignments[rect[0][0]:rect[1][0]+1,rect[0][1]:rect[1][1]+1] = 255

# Convert from mask to training data and train knn
imgLAB = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
trainingData = [imgLAB[:,:,0][assignments != 100],
                imgLAB[:,:,1][assignments != 100],
                imgLAB[:,:,2][assignments != 100]]
trainingData = np.array(trainingData).T
trainingLabels = assignments[assignments != 100]
knn = cv2.KNearest()
knn.train(trainingData.astype(np.float32),trainingLabels.astype(np.float32))

# Run knn and update assignments with results
unknownData = [imgLAB[:,:,0][assignments == 100],
               imgLAB[:,:,1][assignments == 100],
               imgLAB[:,:,2][assignments == 100]]
unknownData = np.array(unknownData).T
ret, newLabels, neighbours, dists = knn.find_nearest(unknownData.astype(np.float32), 2) # TODO arg
assignments[assignments == 100] = newLabels.astype(np.uint8).flatten()

# Perform opening and closing
kernel = np.ones((k,k),np.uint8)
assignments = cv2.morphologyEx(assignments, cv2.MORPH_OPEN, kernel)
assignments = cv2.morphologyEx(assignments, cv2.MORPH_CLOSE, kernel)

# Find largest connected component, maybe specify number of components (TODO)
# Display

# Create masked image
# mask = cv2.threshold(assignments, 100, 255, cv2.THRESH_BINARY)[1]
imageMask = cv2.bitwise_and(img, img, mask=assignments)
imgBgMask = cv2.bitwise_and(bgImg, bgImg, mask=cv2.bitwise_not(assignments))
finalImg = imageMask + imgBgMask

# Display
cv2.imwrite('output/ImageFinal.jpg', finalImg)
cv2.imshow('Final Image', finalImg)
cv2.waitKey(0)

cv2.destroyAllWindows()