import numpy as np
import cv2

BELT_Y_THRESHOLD = 390
OUTPUT_SIZE = 320

def garbage_extract_no_preprocess(img, showCompare = False):
    """
        Extracts the garbage from image
        img: A 4:3 image
        Returns the object's image with size 320x320, black borders, scale to fit

        if showCompare is True, returns the image with annotations without clipping
    """

    # Flip if needed
    if img.shape[0] < img.shape[1]:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) 

    # Resize
    img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)

    output = img.copy()

    # Chroma Keying (blue)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    maskKey = cv2.inRange(hsv, (90,100,130),(110,255,255))
    maskKey = cv2.bitwise_not(maskKey)
    output = cv2.bitwise_and(output, output, mask = maskKey)

    # cv2.imshow('Output', hsv)
    # cv2.waitKey(0)
    
    # Find upper borders
    gray = cv2.cvtColor(output, cv2.COLOR_BGR2GRAY)
    _, bw = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(bw[:BELT_Y_THRESHOLD - 10, :], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Fliter small dots
    fil = list(filter(lambda x: cv2.contourArea(x) > 300, contours))

    cv2.drawContours(output, fil, -1, 255, 3)
    # cv2.imshow('Output', output)
    # cv2.waitKey(0)

    x,y,w,h = cv2.boundingRect(np.vstack(fil)) if len(fil) != 0 else (0, 0, 390, 520)

    # Delete left & right
    output[:, :x, :] = 0
    output[:, x+w: , :] = 0

    # Object border
    rect = [x, y, x+w, 500]

    # grabCut
    mask = np.zeros(img.shape[:2],np.uint8)
    bgdModel = np.zeros((1,65),np.float64)
    fgdModel = np.zeros((1,65),np.float64)
    cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
    output = output*mask2[:,:,np.newaxis]

    if showCompare:
        output = cv2.rectangle(output, rect[:2], rect[2:], (0,255,0), 3)
        return output

    # Clip
    # output = output[rect[1]:rect[3], rect[0]:rect[2], :]
    print(rect[1], rect[3], rect[0], rect[2])
    h = rect[3] - rect[1]
    w = rect[2] - rect[0]
    if h > w:
        rect[2] = min(rect[2] + (h-w)//2, output.shape[0]-1)
        rect[0] = max(rect[0] - (h-w)//2, 0)
    else:
        rect[3] = min(rect[3] + (w-h)//2, output.shape[1]-1)
        rect[1] = max(rect[1] - (w-h)//2, 0)
    print(rect[1], rect[3], rect[0], rect[2])
    output = img[rect[1]:rect[3], rect[0]:rect[2], :]
    
    # Resize
    h, w, c = output.shape
    scale = OUTPUT_SIZE/w if w > h else OUTPUT_SIZE/h
    output = cv2.resize(output, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)

    delta_w = OUTPUT_SIZE - output.shape[1]
    delta_h = OUTPUT_SIZE - output.shape[0]
    top, bottom = delta_h//2, delta_h-(delta_h//2)
    left, right = delta_w//2, delta_w-(delta_w//2)

    color = [0, 0, 0]
    output = cv2.copyMakeBorder(output, top, bottom, left, right, cv2.BORDER_CONSTANT,
        value=color)

    return output

def garbage_extract(img, showCompare = False):
    """
        Extracts the garbage from image
        img: A 4:3 image
        Returns the object's image with size 320x320, black borders, scale to fit

        if showCompare is True, returns the image with annotations without clipping
    """

    # Flip if needed
    if img.shape[0] < img.shape[1]:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) 

    # Resize
    img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)

    output = img.copy()

    # Chroma Keying (blue)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    maskKey = cv2.inRange(hsv, (90,100,130),(110,255,255))
    maskKey = cv2.bitwise_not(maskKey)
    output = cv2.bitwise_and(output, output, mask = maskKey)

    # cv2.imshow('Output', hsv)
    # cv2.waitKey(0)
    
    # Find upper borders
    gray = cv2.cvtColor(output, cv2.COLOR_BGR2GRAY)
    _, bw = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
    contours, hierarchy = cv2.findContours(bw[:BELT_Y_THRESHOLD - 10, :], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Fliter small dots
    fil = list(filter(lambda x: cv2.contourArea(x) > 300, contours))

    # cv2.drawContours(output, fil, -1, 255, 3)
    # cv2.imshow('Output', output)
    # cv2.waitKey(0)

    
    x,y,w,h = cv2.boundingRect(np.vstack(fil))

    # Delete left & right
    output[:, :x, :] = 0
    output[:, x+w: , :] = 0

    # Object border
    rect = (x, y, x+w, 500)

    # grabCut
    mask = np.zeros(img.shape[:2],np.uint8)
    bgdModel = np.zeros((1,65),np.float64)
    fgdModel = np.zeros((1,65),np.float64)
    cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
    output = output*mask2[:,:,np.newaxis]

    if showCompare:
        output = cv2.rectangle(output, rect[:2], rect[2:], (0,255,0), 3)
        return output

    # Clip
    output = output[rect[1]:rect[3], rect[0]:rect[2], :]
    
    # Resize
    h, w, c = output.shape
    scale = OUTPUT_SIZE/w if w > h else OUTPUT_SIZE/h
    output = cv2.resize(output, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)

    delta_w = OUTPUT_SIZE - output.shape[1]
    delta_h = OUTPUT_SIZE - output.shape[0]
    top, bottom = delta_h//2, delta_h-(delta_h//2)
    left, right = delta_w//2, delta_w-(delta_w//2)

    color = [0, 0, 0]
    output = cv2.copyMakeBorder(output, top, bottom, left, right, cv2.BORDER_CONSTANT,
        value=color)

    return output

    
if __name__ == "__main__":
    img = cv2.imread('Camera/paper/paper_leafyellow_14.jpg')

    # compare = garbage_extract(img, True)
    output = garbage_extract(img)

    # Flip if needed
    if img.shape[0] < img.shape[1]:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) 

    # Resize
    img = cv2.resize(img, (390, 520), interpolation=cv2.INTER_AREA)
    
    # cv2.imshow('Compare', np.hstack([compare, img]))
    cv2.imshow('Output', output)

    cv2.waitKey(0)
    cv2.destroyAllWindows()