# from skimage.color import rgb2yuv
# import cv2
import numpy as np
import torch

# from fast_slic.avx2 import SlicAvx2 as Slic
# from skimage.segmentation import slic as sk_slic

from models.UnSAMFlow.utils.warp_utils import flow_warp


# def run_slic_pt(img_batch, n_seg=200, compact=10, rd_select=(8, 16), fast=True):  # Nx1xHxW
#     """

#     :param img: Nx3xHxW 0~1 float32
#     :param n_seg:
#     :param compact:
#     :return: Nx1xHxW float32
#     """
#     B = img_batch.size(0)
#     dtype = img_batch.type()
#     img_batch = np.split(
#         img_batch.detach().cpu().numpy().transpose([0, 2, 3, 1]), B, axis=0)
#     out = []
#     if fast:
#         fast_slic = Slic(num_components=n_seg, compactness=compact, min_size_factor=0.8)
#     for img in img_batch:
#         img = np.copy((img * 255).squeeze(0).astype(np.uint8), order='C')
#         if fast:
#             img = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
#             seg = fast_slic.iterate(img)
#         else:
#             seg = sk_slic(img, n_segments=200, compactness=10)

#         if rd_select is not None:
#             n_select = np.random.randint(rd_select[0], rd_select[1])
#             select_list = np.random.choice(range(0, np.max(seg) + 1), n_select,
#                                            replace=False)

#             seg = np.bitwise_or.reduce([seg == seg_id for seg_id in select_list])
#         out.append(seg)
#     x_out = torch.tensor(np.stack(out)).type(dtype).unsqueeze(1)
#     return x_out


def random_crop(img, full_segs, flow, occ_mask, crop_sz):
    """

    :param img: Nx6xHxW
    :param flows: n * [Nx2xHxW]
    :param occ_masks: n * [Nx1xHxW]
    :param crop_sz:
    :return:
    """
    _, _, h, w = img.size()
    c_h, c_w = crop_sz

    if c_h == h and c_w == w:
        return img, flow, occ_mask

    x1 = np.random.randint(0, w - c_w)
    y1 = np.random.randint(0, h - c_h)
    img = img[:, :, y1 : y1 + c_h, x1 : x1 + c_w]
    full_segs = full_segs[:, :, y1 : y1 + c_h, x1 : x1 + c_w]
    flow = flow[:, :, y1 : y1 + c_h, x1 : x1 + c_w]
    occ_mask = occ_mask[:, :, y1 : y1 + c_h, x1 : x1 + c_w]

    return img, full_segs, flow, occ_mask


# def semantic_connected_components(
#     semseg, class_indices, width_range=(30, 200), height_range=(30, 100)
# ):
#     """
#     Input:
#         semsegs: Onehot semantic segmentations of size [c, H, W]
#         class_indices: A list of the indices of the classes of interest.
#                        For example, [car_idx] or [sign_idx, pole_idx, traffic_light_idx]
#         width_range, height_range: The width and height ranges for the objects of interest.
#     Output:
#         list of masks for cars in the size range
#     """

#     curr_sem = semseg[class_indices].sum(dim=0)
#     curr_sem = (curr_sem[:, :, None] * 255).numpy().astype(np.uint8)
#     num_labels, labels = cv2.connectedComponents(curr_sem)

#     sem_list = []
#     # 0 is background, so ignore them.
#     for i in range(1, num_labels):
#         curr_obj = labels == i
#         hs, ws = np.where(curr_obj)
#         h_len, w_len = np.max(hs) - np.min(hs), np.max(ws) - np.min(ws)
#         if (height_range[0] < h_len < height_range[1]) and (
#             width_range[0] < w_len < width_range[1]
#         ):
#             if (
#                 curr_obj.sum() / ((h_len + 1) * (w_len + 1)) > 0.6
#             ):  # filter some wrong car estimates or largely occluded cars
#                 sem_list.append(curr_obj)

#     return sem_list


# def find_semantic_group(semseg, class_indices, win_width=200):
#     curr_sem = semseg[class_indices].sum(dim=0).numpy()
#     freq = curr_sem.mean(axis=0)  # 1d frequency
#     freq_win = np.convolve(
#         freq, np.ones(win_width) / win_width, mode="valid"
#     )  # find the most frequent window

#     if max(freq_win) > 0.1:
#         # optimal window: [left, right)
#         left = np.argmax(freq_win)
#         right = left + win_width
#         curr_sem[:, :left] = 0
#         curr_sem[:, right:] = 0
#         return curr_sem

#     else:
#         return None


def add_fake_object(input_dict):

    # prepare input
    img1_ot = input_dict["img1_tgt"]
    img2_ot = input_dict["img2_tgt"]
    full_seg1_ot = input_dict["full_seg1_st"]
    full_seg2_ot = input_dict["full_seg2_st"]
    flow_ot = input_dict["flow_tgt"]
    noc_ot = input_dict["noc_tgt"]

    img = input_dict["img_src"]
    obj_mask = input_dict["obj_mask"]
    motion = input_dict["motion"][:, :, None, None]

    b, _, h, w = img1_ot.shape
    N1 = full_seg1_ot.max()
    N2 = full_seg2_ot.max()

    # add object to frame 1
    img1_ot = obj_mask * img + (1 - obj_mask) * img1_ot
    full_seg1_ot = obj_mask * (N1 + 1) + (1 - obj_mask) * full_seg1_ot

    # add object to frame 2
    new_obj_mask = flow_warp(obj_mask, -motion.repeat(1, 1, h, w), pad="zeros")
    new_img = flow_warp(img, -motion.repeat(1, 1, h, w), pad="border")
    img2_ot = new_obj_mask * new_img + (1 - new_obj_mask) * img2_ot
    full_seg2_ot = new_obj_mask * (N2 + 1) + (1 - new_obj_mask) * full_seg2_ot

    # change flow
    flow_ot = obj_mask * motion + (1 - obj_mask) * flow_ot
    noc_ot = torch.max(noc_ot, obj_mask)  # where we are confident about flow_ot

    return img1_ot, img2_ot, full_seg1_ot, full_seg2_ot, flow_ot, noc_ot, new_obj_mask