import os
import json
import base64
import argparse
from glob import glob
from pathlib import Path

def arg_directory(path):
    if os.path.isdir(path):
        return path
    else:
        raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')

def create_directory(directory):
    try:
        os.makedirs(os.path.join(directory,'labels'), exist_ok=True)
    except Exception as e:
        print(f'`{path}` has been created')
    return (directory)


def convert_labelme_to_yolo(labelme_annotation_path, yolo_directory):
    # Load LabelMe annotation
    image_id = Path(labelme_annotation_path).stem
    with open(labelme_annotation_path, 'r') as labelme_annotation_file:
        labelme_annotation = json.load(labelme_annotation_file)

    # YOLO annotation and image paths
    yolo_annotation_path = os.path.join(yolo_directory, 'labels', f'{image_id}.txt')
    yolo_image_path = os.path.join(yolo_directory, 'images/all', f'{image_id}.jpg')

    with open(yolo_annotation_path, 'w') as yolo_annotation_file:
        yolo_image_data = base64.b64decode(labelme_annotation['imageData'])

        # Write YOLO image
        with open(yolo_image_path, 'wb') as yolo_image_file:
            yolo_image_file.write(yolo_image_data)

        # Write YOLO image annotation
        for shape in labelme_annotation['shapes']:
            if shape['shape_type'] != 'rectangle':
                print(f'Invalid type `{shape["shape_type"]}` in annotation `{labelme_annotation_path}`')
                continue

            label = shape['label']

            # shape['points'] format : [[x1,y1],[x2,y2]...] #
            scale_width = 1.0 / labelme_annotation['imageWidth']
            scale_height = 1.0 / labelme_annotation['imageHeight']
            width = abs(shape['points'][1][0] - shape['points'][0][0]) * scale_width
            height = abs(shape['points'][1][1] - shape['points'][0][1]) * scale_height

            x = min(shape['points'][0][0], shape['points'][1][0]) * scale_width + width / 2
            y = min(shape['points'][0][1], shape['points'][1][1]) * scale_height + height / 2
            if x+width/2 > 1 or y+height/2>1:
                print(f'Error with bounding box values over 1 in file {yolo_image_file}')
            annotation_line = f'{label} {x} {y} {width} {height}\n'
            yolo_annotation_file.write(annotation_line)

def main(args):
    yolo_names = set()
    create_directory(os.path.join(args.directory, 'images'))
    create_directory(os.path.join(args.directory, 'images/all'))
    create_directory(os.path.join(args.directory, 'labels'))
    for labelme_annotation_path in glob(f'{args.path_to_data}/*.json'):
        convert_labelme_to_yolo(labelme_annotation_path, args.directory)

        with open(labelme_annotation_path, 'r') as labelme_annotation_file:
            labelme_annotation = json.load(labelme_annotation_file)

            for shape in labelme_annotation['shapes']:
                yolo_names.add(shape['label'])

    # Write YOLO names
    yolo_names_path = os.path.join(args.directory, 'custom.names')
    with open(yolo_names_path, 'w') as yolo_names_file:
        yolo_names_file.write('\n'.join(yolo_names))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Convert LabelMe annotations to YOLO compatible')
    parser.add_argument('-p', '--path_to_data', type=arg_directory, help='Path to LabelMe annotations', required=True)
    parser.add_argument('-d', '--directory', type=arg_directory, help='Directory to which YOLO annotations will be stored', required=True)
    args = parser.parse_args()

    main(args)