import os
import json
import numpy as np
import copy
from PIL import Image
from lmdeploy.vl.constants import IMAGE_TOKEN
from pycocotools import mask as mask_utils
import torch.nn.functional as F
import torch

import cv2

from projects.mllm_labeling.image_blending_fn import contour_rendering

def get_masks_from_anns(annotations):
    captions = []
    masks = []
    for annotation in annotations:
        rle = annotation['object_anno']['segmentation']
        mask = mask_utils.decode(rle)
        masks.append(mask)
        captions.append(annotation['caption'])
    masks = np.stack(masks, axis=0)
    return masks, captions


image_folder = './1215_demos/mask_outs/out/'
save_dir = './1215_demos/overall_demos/'

image_files = []
for file_name in os.listdir(image_folder):
    if 'out' not in file_name and '.jpg' in file_name:
        image_files.append(file_name)

annotation_folder = '1215_demos/mllm_object_cap/'
anno_files = os.listdir(annotation_folder)
annotations = []
for anno_file in anno_files:
    with open(os.path.join(annotation_folder, anno_file), 'r') as f:
        annotations += json.load(f)

image2anno_dict = {}
for annotation in annotations:
    image_id = annotation['image_id']
    if image_id not in image2anno_dict.keys():
        image2anno_dict[image_id] = [annotation]
    else:
        image2anno_dict[image_id].append(annotation)

for i, image_name in enumerate(image2anno_dict.keys()):
    txt_strs = ''
    print('====================================================')
    image_path = os.path.join(image_folder, image_name)
    image = Image.open(image_path).convert('RGB')

    image_annotations = image2anno_dict[image_name]
    masks, captions = get_masks_from_anns(image_annotations)

    anno_ids = []

    for i_obj, caption in enumerate(captions):
        txt_strs += '\n\n' + f"{i_obj}: " + caption + '\n\n'
        anno_ids.append(i_obj)
        print(caption)
        print('+++++++++++++++++++++++++++')

    image_shape = image.size
    masks = torch.Tensor(masks).unsqueeze(0)
    masks = F.interpolate(
        masks,
        size=(image_shape[1], image_shape[0]),
        mode='nearest').squeeze(0)
    masks = masks.numpy().astype(np.uint8)

    image = np.array(image)
    contour_rendering(image, masks, mask_ids=anno_ids)

    image = Image.fromarray(image)
    image.save(os.path.join(save_dir, f"{i}.png"))
    with open(os.path.join(save_dir, f"{i}.txt"), 'w') as f:
        f.write(txt_strs)