import os
import json
import numpy as np
import cv2
import shutil
from tqdm import tqdm


def labelme2mask_single_img(img_path, labelme_json_path, class_info):
    '''
    Convert a single image's LabelMe annotation to a mask.
    '''
    img_bgr = cv2.imread(img_path)
    img_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8)  # Create an empty mask image (0 - background)

    with open(labelme_json_path, 'r', encoding='utf-8') as f:
        labelme = json.load(f)

    for one_class in class_info:  # Iterate over each class in class_info
        for each in labelme['shapes']:  # Iterate over all shapes in the annotation
            if each['label'] == one_class['label']:
                if one_class['type'] == 'polygon':  # Handle polygon annotation
                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]  # Ensure correct shape
                    img_mask = cv2.fillPoly(img_mask, points, color=one_class['color'])
                elif one_class['type'] == 'line' or one_class['type'] == 'linestrip':  # Handle line annotation
                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
                    img_mask = cv2.polylines(img_mask, points, isClosed=False, color=one_class['color'],
                                             thickness=one_class.get('thickness', 1))
                elif one_class['type'] == 'circle':  # Handle circle annotation
                    points = np.array(each['points'], dtype=np.int32)
                    center_x, center_y = points[0][0], points[0][1]
                    edge_x, edge_y = points[1][0], points[1][1]
                    radius = int(np.linalg.norm([center_x - edge_x, center_y - edge_y]))
                    img_mask = cv2.circle(img_mask, (center_x, center_y), radius, one_class['color'], -1)
                else:
                    print('Unknown annotation type:', one_class['type'])

    return img_mask


def convert_labelme_to_mask(Dataset_Path):
    '''
    Convert all LabelMe annotations in the dataset to mask images.
    '''
    # Dataset directories
    img_dir = os.path.join(Dataset_Path, 'images')
    ann_dir = os.path.join(Dataset_Path, 'labelme_jsons')

    # Class information for mask generation
    class_info = [
        {'label': 'panicle', 'type': 'polygon', 'color': 1}
    ]

    # Create target directories
    images_target_dir = os.path.join(Dataset_Path, 'img_dir')
    ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')

    # Create target directories if they do not exist
    os.makedirs(images_target_dir, exist_ok=True)
    os.makedirs(ann_target_dir, exist_ok=True)

    # Process each image in the images directory
    for img_name in tqdm(os.listdir(img_dir), desc="Converting images to masks"):
        try:
            img_path = os.path.join(img_dir, img_name)
            labelme_json_path = os.path.join(ann_dir, f'{os.path.splitext(img_name)[0]}.json')

            if os.path.exists(labelme_json_path):
                # Convert LabelMe annotations to mask
                img_mask = labelme2mask_single_img(img_path, labelme_json_path, class_info)

                # Save the mask to the target directory
                mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')
                cv2.imwrite(mask_path, img_mask)

                # Move the image to the target directory
                shutil.move(img_path, os.path.join(images_target_dir, img_name))
            else:
                print(f"Annotation file missing for {img_name}")

        except Exception as e:
            print(f"Failed to convert {img_name}: {e}")

    # Optionally remove the original directories if they are empty
    shutil.rmtree(img_dir, ignore_errors=True)
    shutil.rmtree(ann_dir, ignore_errors=True)

    print("Conversion completed.")


if __name__ == '__main__':
    Dataset_Path = 'CVRP'  # Update this to the path of your dataset
    convert_labelme_to_mask(Dataset_Path)