#!/usr/bin/env python

from __future__ import print_function

import argparse
import glob
import os
import os.path as osp
import sys

import imgviz
import labelme

try:
    import lxml.builder
    import lxml.etree
except ImportError:
    print("Please install lxml:\n\n    pip install lxml\n")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", help="input annotated directory")
    parser.add_argument("--output_dir", help="output dataset directory")
    parser.add_argument("--labels", help="labels file", required=True)
    parser.add_argument(
        "--noviz", help="no visualization", action="store_true"
    )
    args = parser.parse_args()

    if not osp.exists(args.output_dir):
        os.makedirs(args.output_dir)
    # os.makedirs(args.output_dir)
    # os.makedirs(osp.join(args.output_dir, "JPEGImages"))
    # os.makedirs(osp.join(args.output_dir, "Annotations"))
    if not args.noviz:
        os.makedirs(osp.join(args.output_dir, "AnnotationsVisualization"))
    print("Creating dataset:", args.output_dir)

    class_names = []
    class_name_to_id = {}
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i
        class_name = line.strip()
        class_name_to_id[class_name] = class_id
        class_names.append(class_name)
    class_names = tuple(class_names)
    print("class_names:", class_names)
    out_class_names_file = osp.join(args.output_dir, "class_names.txt")
    # with open(out_class_names_file, "w") as f:
    #     f.writelines("\n".join(class_names))
    # print("Saved class_names:", out_class_names_file)
    
    out_name = os.path.basename(args.output_dir) + ".xml"
    out_xml_file = osp.join(args.output_dir, out_name)
    root_maker = lxml.builder.ElementMaker()
    root = root_maker.data()
    for filename in glob.glob(osp.join(args.input_dir, "*.json")):
        print("Generating dataset from:", filename)

        label_file = labelme.LabelFile(filename=filename)

        base = osp.splitext(osp.basename(filename))[0]
        out_img_file = osp.join(args.output_dir, "JPEGImages", base + ".jpg")
        if not args.noviz:
            out_viz_file = osp.join(
                args.output_dir, "AnnotationsVisualization", base + ".jpg"
            )

        img = labelme.utils.img_data_to_arr(label_file.imageData)
        # imgviz.io.imsave(out_img_file, img)

        maker = lxml.builder.ElementMaker()
        xml = maker.frame(
            maker.filename(base + ".jpg"),
        )

        points = []
        labels = []
        for shape in label_file.shapes:
            if shape["shape_type"] != "point":
                print(
                    "Skipping shape: label={label}, "
                    "shape_type={shape_type}".format(**shape)
                )
                continue

            class_name = shape["label"]
            class_id = class_names.index(class_name)

            point = shape["points"]

            points.append(point[0])
            labels.append(class_id)

            xml.append(
                maker.point(
                    name=shape["label"],
                    x=str(point[0][0]),
                    y=str(point[0][1])
                )
            )
           
        root.append(xml)
        if not args.noviz:
            captions = [class_names[label] for label in labels]
            viz = imgviz.instances2rgb(
                image=img,
                labels=labels,
                bboxes=points,
                captions=captions,
                font_size=15,
            )
            imgviz.io.imsave(out_viz_file, viz)

    with open(out_xml_file, "wb") as f:
        f.write(lxml.etree.tostring(root, pretty_print=True))


if __name__ == "__main__":
    main()