import json
import os
from PIL import Image
import numpy as np
from pycocotools.mask import encode, decode, frPyObjects
from tqdm import tqdm
import copy
from natsort import natsorted

if __name__ == '__main__':
    root_path = '/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_dataset_mugs/test'
    save_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_dataset_mugs/handal_datasets_mugs_test.json"
    val_set = os.listdir(root_path)
    new_img_id = 0
    handal_dataset = []
    for val_name in tqdm(val_set):
        vid_path = os.path.join(root_path, val_name)
        img_path = os.path.join(vid_path, "rgb")
        anno_path = os.path.join(vid_path, "mask")
        frame_idx = natsorted(os.listdir(img_path))
        frame_idx  = [f.split(".")[0] for f in frame_idx]
        video_len = len(frame_idx)
        for i,idx in enumerate(frame_idx):
            if i+100 > video_len-1:
                break
            target_idx = frame_idx[i+100]

            first_frame_annotation_path = os.path.join(anno_path, idx+"_000000.png")
            first_frame_annotation_relpath = os.path.relpath(first_frame_annotation_path, root_path)

            first_frame_img_path = os.path.join(img_path, idx+".jpg")
            first_frame_img_relpath = os.path.relpath(first_frame_img_path, root_path)

            first_frame_annotation_img = Image.open(first_frame_annotation_path)
            first_frame_annotation = np.array(first_frame_annotation_img)
            height, width = first_frame_annotation.shape
            unique_instances = np.unique(first_frame_annotation)
            unique_instances = unique_instances[unique_instances != 0]
            coco_format_annotations = []
            # for semi-supervised VOS, we use first frame's GT for input
            for instance_value in unique_instances:
                binary_mask = (first_frame_annotation == instance_value).astype(np.uint8)
                segmentation = encode(np.asfortranarray(binary_mask))
                segmentation = {
                    'counts': segmentation['counts'].decode('ascii'),
                    'size': segmentation['size'],
                }
                area = binary_mask.sum().astype(float)
                coco_format_annotations.append(
                    {
                        'segmentation': segmentation,
                        'area': area,
                        'category_id': instance_value.astype(float),
                    }
                )

            sample_img_path = os.path.join(img_path, target_idx+".jpg")
            sample_img_relpath = os.path.relpath(sample_img_path, root_path)
            image_info = {
                'file_name': sample_img_relpath,
                'height': height,
                'width': width,
            }
            sample_annotation_path = os.path.join(anno_path, target_idx+"_000000.png")
            sample_annotation = np.array(Image.open(sample_annotation_path))

            sample_unique_instances = np.unique(sample_annotation)
            sample_unique_instances = sample_unique_instances[sample_unique_instances != 0]
            anns = []
            for instance_value in sample_unique_instances:
                assert instance_value in unique_instances, 'Found new target not in the first frame'
                binary_mask = (sample_annotation == instance_value).astype(np.uint8)
                segmentation = encode(np.asfortranarray(binary_mask))
                segmentation = {
                    'counts': segmentation['counts'].decode('ascii'),
                    'size': segmentation['size'],
                }
                area = binary_mask.sum().astype(float)
                anns.append(
                    {
                        'segmentation': segmentation,
                        'area': area,
                        'category_id': instance_value.astype(float),
                    }
                )
            first_frame_anns = copy.deepcopy(coco_format_annotations)
            if len(anns) < len(first_frame_anns):
                first_frame_anns = [ann for ann in first_frame_anns if ann['category_id'] in sample_unique_instances]
            assert len(anns) == len(first_frame_anns)
            sample = {
                'image': sample_img_relpath,
                'image_info': image_info,
                'anns': anns,
                'first_frame_image': first_frame_img_relpath,
                'first_frame_anns': first_frame_anns,
                'new_img_id': new_img_id,
                'video_name': val_name,
            }
            handal_dataset.append(sample)
            new_img_id += 1
    
   
    with open(save_path, 'w') as f:
        json.dump(handal_dataset, f)
    print(f'Save at {save_path}. Total sample: {len(handal_dataset)}')