import numpy as np

from copy import deepcopy

from tqdm import tqdm

import cv2

from multiprocessing import Pool, cpu_count


def parse_joint(list1, list2, parse_key):

    list1_ids = list(map(parse_key, list1))
    list2_ids = list(map(parse_key, list2))

    valid_idx = []

    for i in range(len(list1_ids)):

        if list1_ids[i] in list2_ids:

            valid_idx.append(i)

    return valid_idx


def kmeans_segmentation_step(image, n_clusters=2, niter=20):

    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)

    image1d = image.reshape(-1, 3)

    ret, label, center = cv2.kmeans(image1d,
                                    n_clusters,
                                    None,
                                    criteria,
                                    niter,
                                    cv2.KMEANS_PP_CENTERS)

    mask = label.reshape(image.shape[:2])

    return mask


def kmeans_binary_segmentation(images):

    with Pool(cpu_count()) as pool:

        masks = list(tqdm(pool.imap(kmeans_segmentation_step, images), total=len(images)))

    return np.array(masks)
