import numpy as np

import colorsys

import random

import cv2

def get_n_hls_colors(num):

    hls_colors = []

    i = 0

    step = 360.0 / num

    while i < 360:

        h = i

        s = 90 + random.random() * 10

        l = 50 + random.random() * 10

        _hlsc = [h / 360.0, l / 100.0, s / 100.0]

        hls_colors.append(_hlsc)

        i += step

    return hls_colors

def ncolors(num):

    rgb_colors = []

    if num < 1:

        return np.array(rgb_colors)

    hls_colors = get_n_hls_colors(num)

    for hlsc in hls_colors:

        _r, _g, _b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])

        rgb_colors.append([_r, _g, _b])

        # r, g, b = [int(x * 255.0) for x in (_r, _g, _b)]

        # rgb_colors.append([r, g, b])

    return np.array(rgb_colors)

def random_colors(N, bright=True):

    """

    Generate random colors.

    To get visually distinct colors, generate them in HSV space then

    convert to RGB.

    """

    brightness = 1.0 if bright else 0.7

    hsv = [(i / N, 1, brightness) for i in range(N)]

    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))

    # random.shuffle(colors)

    return colors

def mask2png(mask, file_name=None, suffix="png"):

    """ mask: (w, h)

        img_rgb: (w, h, rgb)

    """

    nums = np.unique(mask)[1:]

    if len(nums) < 1:

        colors = np.array([[0,0,0]])

    else:

        # colors = ncolors(len(nums))

        colors = (np.array(random_colors(len(nums))) * 255).astype(int)

        colors = np.insert(colors, 0, [0,0,0], 0)

    # 保证mask中的值为1-N连续

    mask_ordered = np.zeros_like(mask)

    for cnt, l in enumerate(nums, 1):

        mask_ordered[mask==l] = cnt

    im_rgb = colors[mask_ordered.astype(int)]

    if file_name is not None:

        cv2.imwrite(file_name+"."+suffix, im_rgb[:, :, ::-1])

    return im_rgb

def apply_mask(image, mask, color, alpha=0.5, scale=1):

    """Apply the given mask to the image.

    """

    for c in range(3):

        image[:, :, c] = np.where(mask == 1,

                                  image[:, :, c] *

                                  (1 - alpha) + alpha * color[c] * scale,

                                  image[:, :, c])

    return image

def img_mask_png(image, mask, file_name=None, alpha=0.5, suffix="png"):

    """ mask: (h, w)

        image: (h, w, rgb)

    """

    nums = np.unique(mask)[1:]

    if len(nums) < 1:

        colors = np.array([[0,0,0]])

    else:

        colors = ncolors(len(nums))

        colors = np.insert(colors, 0, [0,0,0], 0)

    # 保证mask中的值为1-N连续

    mask_ordered = np.zeros_like(mask)

    for cnt, l in enumerate(nums, 1):

        mask_ordered[mask==l] = cnt

    # mask_rgb = colors[mask_ordered.astype(int)]

    mix_im = image.copy()

    for i in np.unique(mask_ordered)[1:]:

        mix_im = apply_mask(mix_im, mask_ordered==i, colors[int(i)], alpha=alpha, scale=255)

    if file_name is not None:

        cv2.imwrite(file_name+"."+suffix, mix_im[:, :, ::-1])

    return mix_im

def _find_contour(mask):

    # _, contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 顶点

    _, contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    cont = np.zeros_like(mask)

    for contour in contours:

        cont[contour[:,:,1], contour[:,:,0]] = 1

    return cont

def masks_to_contours(masks):

    # 包含多个区域

    nums = np.unique(masks)[1:]

    cont_mask = np.zeros_like(masks)

    for i in nums:

        cont_mask += _find_contour(masks==i)

    return (cont_mask>0).astype(int)

def batchToColorImg(batch, minv=None, maxv=None, scale=255.):

    """ batch: (N, H, W, C)

    """

    if batch.ndim == 3:

        N, H, W = batch.shape

    elif batch.ndim == 4:

        N, H, W, _ = batch.shape

    colorImg = np.zeros(shape=(N, H, W, 3))

    for i in range(N):

        if minv is None:

            a = (batch[i] - batch[i].min()) / (batch[i].max() - batch[i].min()) * 255

        else:

            a = (batch[i] - minv) / (maxv - minv) * scale

        a = cv2.applyColorMap(a.astype(np.uint8), cv2.COLORMAP_JET)

        colorImg[i, ...] = a[..., ::-1] / 255.

    return colorImg

if __name__ == "__main__":

    a = np.zeros((100, 100))

    a[0:5, 3:8] = 1

    a[75:85, 85:95] = 2

    # colors = ncolors(2)[::-1]

    colors = np.array(random_colors(2)) * 255

    colors = np.insert(colors, 0, [0, 0, 0], 0)

    b = colors[a.astype(int)].astype(np.uint8)

    import cv2, skimage

    # skimage.io.imsave("test_io.png", b)

    # # cv2.imwrite("test.jpg", b[:, :, ::-1])

    # print()

    # mask2png(a, "test")

    ################################################

    # img_mask_png(b, a, "test")

    #############################################

    # cont_mask = find_contours(a)

    # print()

    # # skimage.io.imsave("test_cont.png", cont_mask) 

    # b[cont_mask>0, :] = [255, 255, 255]

    # skimage.io.imsave("test_cont.png", b)

    gt0 = skimage.io.imread("gt0.png", as_gray=False)

    print()

    gt0[gt0==54] = 0

    # cont_mask = find_contours(gt0==237)  # Array([ 18,  54,  73, 237], dtype=uint8)

    # cont_mask += find_contours(gt0==18)

    # cont_mask += find_contours(gt0==73)

    cont_mask = masks_to_contours(gt0)

    colors = np.array(random_colors(1)) * 255

    colors = np.insert(colors, 0, [0, 0, 0], 0)

    cont_mask = colors[cont_mask.astype(int)].astype(np.uint8)

    skimage.io.imsave("test_cont.png", cont_mask)