import numpy as np
import cv2
import os

def image_resize(img, width = 400, height = 400):
    return cv2.resize(img, (width,height))

def grabcut_rect(image_path, rect_coords, project_name, row_id):
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    dimy = np.shape(img)[0]
    dimx = np.shape(img)[1]
    mask = np.zeros((dimy, dimx), np.uint8)  # initialize empty mask
    bgdModel = fgdModel = np.zeros((1, 65), np.float64)

    x = int(rect_coords['left'])
    y = int(rect_coords['top'])
    h = int(rect_coords['height'])
    w = int(rect_coords['width'])
    rect = (x, y, w, h)

    mask, bgdModel, fgdModel = cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)

    img_file_name = new_file_name(project_name, row_id)

    mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
    img2 = img * mask2[:, :, np.newaxis]

    # to save the image without black background
    tmp = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
    b, g, r = cv2.split(img2)
    rgba = [b, g, r, alpha]
    img2_refined = cv2.merge(rgba, 4)
    cv2.imwrite(img_file_name, img2_refined)

    # send the new file name back
    return img_file_name


def resize_and_write(img_path, img, rect_coords):
    print(rect_coords)
    # img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
    x = int(rect_coords['left'])
    y = int(rect_coords['top'])
    h = int(rect_coords['height'])
    w = int(rect_coords['width'])
    print(rect_coords)
    crop_img = img[y:y + h, x:x + w]
    cv2.imwrite(img_path, crop_img)


# give a unique name to each saved file grabcut_0,1,2,...
# def new_file_name(image_path):
def new_file_name(project_name, row_id):
    # index = image_path.rfind('.')

    file_path = "./static/images/" + project_name + "/" + str(row_id) + "/"
    if (not os.path.isdir(file_path)):
        os.mkdir(file_path)
    new_name = file_path + 'grabcut'
    i = 0
    while os.path.isfile(new_name + str(i) + '.png'):
        i += 1
    return new_name + str(i) + '.png'


def grabcut_drawing(image_path, rect_coords, fg_drawing, bg_drawing, project_name, row_id):
    print('in refine grab cut!')
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    dimy = np.shape(img)[0]
    dimx = np.shape(img)[1]

    mask = np.zeros((dimy, dimx), np.uint8)  # initialize empty mask
    bgdModel = fgdModel = np.zeros((1, 65), np.float64)

    x = int(rect_coords['left'])
    y = int(rect_coords['top'])
    h = int(rect_coords['height'])
    w = int(rect_coords['width'])
    rect = (x, y, w, h)  # make rect to initialize with rect
    cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
    bgdModel = np.zeros((1, 65), np.float64)  # zero out foreground and background models
    fgdModel = np.zeros((1, 65), np.float64)

    # make the drawing on the mask, 
    # for each line in the drawing,
    #  for each point in the line, make a circle with the scaled thickness on the mask    
    # lines = drawing['lines']
    # thickness = drawing['thickness']
    # for line in lines:
    #     path = line['path']
    #     for path_seg in path:
    #         path_seg_type = path_seg[0]
    #         if (path_seg_type == 'Q'):
    #             # print('CIRCLE IN MASK')

    #             x1 = int(path_seg[1])
    #             y1 = int(path_seg[2])
    #             if (x1 > dimx or x1 < 0) or (y1> dimy or y1 < 0) :
    #                 continue
    #             cv2.circle(mask, (x1,y1), int(thickness), 0, -1)
    print("foreground drawling:", fg_drawing)
    print("background drawing:", bg_drawing)
    mask = draw_lines_on_mask(mask, fg_drawing, int(cv2.GC_FGD), dimx, dimy)
    mask = draw_lines_on_mask(mask, bg_drawing, int(cv2.GC_BGD), dimx, dimy)

    mask, bgdModel, fgdModel = cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)
    mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
    img2 = img * mask2[:, :, np.newaxis]

    img_file_name = new_file_name(project_name, row_id)

    # to save the image without black background
    tmp = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY)
    b, g, r = cv2.split(img2)
    rgba = [b, g, r, alpha]
    img2_refined = cv2.merge(rgba, 4)
    cv2.imwrite(img_file_name, img2_refined)

    # send the new file name back
    return img_file_name


def draw_lines_on_mask(mask, drawing, val, dimx, dimy):
    lines = drawing['lines']
    thickness = drawing['thickness']
    for line in lines:
        path = line['path']
        for path_seg in path:
            path_seg_type = path_seg[0]
            if (path_seg_type == 'Q'):
                # print('CIRCLE IN MASK')

                x1 = int(path_seg[1])
                y1 = int(path_seg[2])
                if (x1 > dimx or x1 < 0) or (y1 > dimy or y1 < 0):
                    continue
                cv2.circle(mask, (x1, y1), int(thickness), val, -1)
    return mask
