import argparse
import json
import tqdm
import cv2
import os
import numpy as np
from pycocotools import mask as mask_utils
import random
from PIL import Image
from natsort import natsorted

EVALMODE = "test"
# MAX_OBJ_NUM = 2 # debug

def blend_mask(input_img, binary_mask, alpha=0.5, color="g"):
    if input_img.ndim == 2:
        return input_img
    mask_image = np.zeros(input_img.shape, np.uint8)
    if color == "r":
        mask_image[:, :, 0] = 255
    if color == "g":
        mask_image[:, :, 1] = 255
    if color == "b":
        mask_image[:, :, 2] = 255
    if color == "o":
        mask_image[:, :, 0] = 255
        mask_image[:, :, 1] = 165
        mask_image[:, :, 2] = 0
    if color == "c":
        mask_image[:, :, 0] = 0
        mask_image[:, :, 1] = 255
        mask_image[:, :, 2] = 255
    if color == "p":
        mask_image[:, :, 0] = 128
        mask_image[:, :, 1] = 0
        mask_image[:, :, 2] = 128

    mask_image = mask_image * np.repeat(binary_mask[:, :, np.newaxis], 3, axis=2)
    blend_image = input_img[:, :, :].copy()
    pos_idx = binary_mask > 0
    for ind in range(input_img.ndim):
        ch_img1 = input_img[:, :, ind]
        ch_img2 = mask_image[:, :, ind]
        ch_img3 = blend_image[:, :, ind]
        ch_img3[pos_idx] = alpha * ch_img1[pos_idx] + (1 - alpha) * ch_img2[pos_idx]
        blend_image[:, :, ind] = ch_img3
    return blend_image


def upsample_mask(mask, frame):
    H, W = frame.shape[:2]
    mH, mW = mask.shape[:2]

    if W > H:
        ratio = mW / W
        h = H * ratio
        diff = int((mH - h) // 2)
        if diff == 0:
            mask = mask
        else:
            mask = mask[diff:-diff]
    else:
        ratio = mH / H
        w = W * ratio
        diff = int((mW - w) // 2)
        if diff == 0:
            mask = mask
        else:
            mask = mask[:, diff:-diff]

    mask = cv2.resize(mask, (W, H))
    return mask


def downsample(mask, frame):
    H, W = frame.shape[:2]
    mH, mW = mask.shape[:2]

    mask = cv2.resize(mask, (W, H))
    return mask


#datapath /datasegswap
#inference_path /inference_xmem_ego_last/coco
#output /vis_piano
#--show_gt要加上
if __name__ == "__main__":

    color = ['g', 'r', 'b', 'o', 'c', 'p']
    data_path = "/scratch/yuqian_fu/data_imgs/b511dfed-58f4-4c91-bf0a-f8ce9d47aea9/cam01"
    output_path = "/scratch/yuqian_fu/iccv_vis_imgs/sam2/exo"
    # output_path = "/scratch/yuqian_fu/iccv_vis_imgs/ego2exo"
    mask_path = "/scratch/yuqian_fu/iccv_results_memory/sam2/exo"
    # mask_path = "/scratch/yuqian_fu/iccv_results_memory/ego2exo"
    takes_ids = ["b511dfed-58f4-4c91-bf0a-f8ce9d47aea9"]

    for take_id in tqdm.tqdm(takes_ids):
        prediction_path = os.path.join(mask_path, take_id)
        file_names = natsorted(os.listdir(prediction_path))
        idxs = [f.split(".")[0] for f in file_names]
 
        INSTANCE_ALL = []
        # 先统计该take下所有的物体个数及对应id
        for id in idxs:
            frame_idx = id
            frame = cv2.imread(
                f"{data_path}/{frame_idx}.jpg"
            )
            mask = Image.open(f"{prediction_path}/{frame_idx}.png")
            mask = np.array(mask)
            unique_instances = np.unique(mask)
            unique_instances = unique_instances[unique_instances != 0]
            if len(unique_instances) > 0:
                for instance_value in unique_instances:
                    if instance_value not in INSTANCE_ALL:
                        INSTANCE_ALL.append(instance_value)
        print("INSTANCE_ALL:", INSTANCE_ALL) # debug

        # 开始上色
        for id in idxs:
            frame_idx = id
            frame = cv2.imread(
                f"{data_path}/{frame_idx}.jpg"
            )
            mask = Image.open(f"{prediction_path}/{frame_idx}.png")
            mask = np.array(mask)
            unique_instances = np.unique(mask)
            unique_instances = unique_instances[unique_instances != 0]
            # print("unique_instances:", unique_instances) # debug
            if len(unique_instances) == len(INSTANCE_ALL):
                for i,instance_value in enumerate(unique_instances):
                    binary_mask = (mask == instance_value).astype(np.uint8)
                    binary_mask = cv2.resize(binary_mask, (frame.shape[1], frame.shape[0]))
                    binary_mask = upsample_mask(binary_mask, frame)
                    out = blend_mask(frame, binary_mask, color=color[0])
                    os.makedirs(
                                f"{output_path}/{take_id}/obj_{instance_value}",  #debug
                                exist_ok=True,
                            )
                    cv2.imwrite(
                        f"{output_path}/{take_id}/obj_{instance_value}/{frame_idx}.jpg",
                        out,
                    )
            
            elif len(unique_instances) < len(INSTANCE_ALL):
                for i, instance_value in enumerate(INSTANCE_ALL):
                    if instance_value in unique_instances:
                        binary_mask = (mask == instance_value).astype(np.uint8)
                        binary_mask = cv2.resize(binary_mask, (frame.shape[1], frame.shape[0]))
                        binary_mask = upsample_mask(binary_mask, frame)
                        out = blend_mask(frame, binary_mask, color=color[0])
                    else:
                        zero_mask = np.zeros((frame.shape[0], frame.shape[1]), dtype=np.uint8)
                        out = blend_mask(frame, zero_mask, color=color[0])
                    
                    os.makedirs(
                                f"{output_path}/{take_id}/obj_{instance_value}",  #debug
                                exist_ok=True,
                            )
                    cv2.imwrite(
                        f"{output_path}/{take_id}/obj_{instance_value}/{frame_idx}.jpg",
                        out,
                    )