import numpy as np
import cv2
import skimage.morphology
import skimage.filters.rank


# Aux functions (originally from suzie.images)

def local_entropy(im, kernel_size=5, normalize=True):
    kernel = skimage.morphology.disk(kernel_size)
    if len(im.shape)>2:
        kkernel = np.zeros((kernel.shape[0], kernel.shape[1],im.shape[2]))
        for i in range(im.shape[2]):
            kkernel[:,:,i]=kernel
        kernel = kkernel
    entr_img = skimage.filters.rank.entropy(im, kernel)
    if normalize:
        max_img = np.max(entr_img)
        entr_img = (entr_img * 255 / max_img).astype(np.uint8)
    return entr_img


def calc_dim(contour):
    c_0 = [point[0][0] for point in contour]
    c_1 = [point[0][1] for point in contour]
    return (min(c_0), max(c_0), min(c_1), max(c_1))


def calc_size(dim):
    return (dim[1] - dim[0]) * (dim[3] - dim[2])


def calc_dist(dim1, dim2):
    return None


# -- Main function

def extract_mask(img, filled=True, threshold=135, kernel_size=5):
    entr_img = local_entropy(img, kernel_size=kernel_size)
    _, mask = cv2.threshold(entr_img, threshold, 255, cv2.THRESH_BINARY)
    while(len(np.unique(mask))==1):
        threshold -= 5
        print("Reducing threshold to ",threshold)
        _, mask = cv2.threshold(entr_img, threshold, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

    contours_d = [calc_dim(c) for c in contours]
    contours_sizes = [calc_size(c) for c in contours_d]
    contour_indices = np.argsort(contours_sizes)[::-1]

    # remove artifacts
    fratio = 0
    sratio = 5
    idx = -1
    while fratio < 0.3 or sratio > 5:
        idx += 1
        biggest = contour_indices[idx]
        filled_mask = np.zeros(img.shape, dtype=np.uint8)
        filled_mask = cv2.fillPoly(filled_mask, [contours[biggest]], 255)
        fratio = filled_mask.sum() / 255 / contours_sizes[biggest]
        cdim = contours_d[biggest]
        sratio = (cdim[3] - cdim[2]) / (cdim[1] - cdim[0])
        if sratio < 1: sratio = 1 / sratio
  
    # generating the mask
    filled_mask = np.zeros(img.shape, dtype=np.uint8)
    if filled:
        filled_mask = cv2.fillPoly(filled_mask, [contours[biggest]], 255)
    else:
        filled_mask = cv2.polylines(filled_mask, [contours[biggest]], 30, 255)

    return filled_mask