import argparse
import json
import tqdm
import cv2
import os
import numpy as np
import random
from pycocotools.mask import encode, decode, frPyObjects

EVALMODE = "test"


def fuse_mask(mask_list):
    fused_mask = np.zeros_like(mask_list[0])
    for mask in mask_list:
        fused_mask[mask == 1] = 1
    return fused_mask


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


if __name__ == "__main__":
   

    color = ['g', 'r', 'b', 'o', 'c', 'p']
    data_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL"
    json_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL_results/xmem_handal.json" #debug
    output_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/vis_rondom_check/xmem"
    

    with open(json_path, "r") as fp:
        datas = json.load(fp)

    # datas = random.sample(datas, 10)
    
    
    for data in datas:
        video_name = data["image"].split("/")[0]
        name = data["image"].split("/")[-1]
        frame_idx = name.split(".")[0]
        #target gt
        frame_target = cv2.imread(
                        f"{data_path}/{data['image']}"
                    )


        # for check xmem predicted result
        mask = decode(data["pred_mask"])
        mask = downsample(mask, frame_target)
        out = blend_mask(frame_target, mask, color=color[0])
        os.makedirs(
                        f"{output_path}/xmem", #debug
                        exist_ok=True,
                    )
        cv2.imwrite(
                        f"{output_path}/xmem/{frame_idx}.jpg",  #debug
                        out,
                    )






        # for check handal original json
        # for i,ann in enumerate(data["anns"]):
        #     mask = decode(ann["segmentation"])
        #     mask = downsample(mask, frame_target)
        #     out = blend_mask(frame_target, mask, color=color[0])
        #     os.makedirs(
        #                 f"{output_path}/{video_name}/target_gt", #debug
        #                 exist_ok=True,
        #             )
        #     cv2.imwrite(
        #                 f"{output_path}/{video_name}/target_gt/{frame_idx}.jpg",  #debug
        #                 out,
        #             )
    
        
        # #query gt
        # frame_query = cv2.imread(
        #                 f"{data_path}/{data['first_frame_image']}"
        #             )


        # for i,ann in enumerate(data["first_frame_anns"]):
        #     mask = decode(ann["segmentation"])
        #     mask = downsample(mask, frame_query)
        #     out = blend_mask(frame_query, mask, color=color[0])
        #     os.makedirs(
        #                 f"{output_path}/{video_name}/query_gt",  #debug
        #                 exist_ok=True,
        #             )
        #     cv2.imwrite(
        #                 f"{output_path}/{video_name}/query_gt/{frame_idx}.jpg",  #debug
        #                 out,
        #             )