"""
Usage:
# Create train data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/train -o [PATH_TO_ANNOTATIONS_FOLDER]/train_labels.csv

# Create test data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/test -o [PATH_TO_ANNOTATIONS_FOLDER]/test_labels.csv
"""

import os
import glob
import pandas as pd
import argparse
import xml.etree.ElementTree as ET
from PIL import Image


def xml_to_csv(path):
    """Iterates through all .xml files (generated by labelImg) in a given directory and combines them in a single Pandas datagrame.

    Parameters:
    ----------
    path : {str}
        The path containing the .xml files
    Returns
    -------
    Pandas DataFrame
        The produced dataframe
    """
    classes_names = []
    txt_list = []
    for txt_file in glob.glob(path + "/*.txt"):
        #tree = ET.parse(txt_file)
        #root = tree.getroot()
        
            classes_names.append("fork")
            #fname = os.path.splitext(os.path.basename(txt_file))[0]
            fname = os.path.splitext(txt_file)[0]
            print(fname)
            #fname = fname[:-1]
            #print(fname)
            imgFile =  fname + ".jpg"
            image = Image.open(imgFile)
            width, height = image.size
            
            with open(txt_file) as fp:
                cnt = 0
                for line in fp:
                    values = str.split(line)
                    minx = int(float(values[1]) * width - float(values[3])*width/2)
                    miny = int(float(values[2]) * height - float(values[4])*height/2)

                    maxx = int(float(values[1]) * width + float(values[3])*width/2)
                    maxy = int(float(values[2]) * height + float(values[4])*height/2)
                   
                    value = (
                        # root.find("filename").text,
                        # int(root.find("size")[0].text),
                        # int(root.find("size")[1].text),
                        # member[0].text,
                        # int(member[4][0].text),
                        # int(member[4][1].text),
                        # int(member[4][2].text),
                        # int(member[4][3].text),
                        imgFile,
                        width,
                        height,
                        "fork",
                        minx,
                        miny,
                        maxx,
                        maxy,

                        
                    )
                    txt_list.append(value)
    column_name = [
        "filename",
        "width",
        "height",
        "class",
        "xmin",
        "ymin",
        "xmax",
        "ymax",
    ]
    xml_df = pd.DataFrame(txt_list, columns=column_name)
    classes_names = list(set(classes_names))
    classes_names.sort()
    return xml_df, classes_names


def main():
    # Initiate argument parser
    parser = argparse.ArgumentParser(
        description="Sample TensorFlow XML-to-CSV converter"
    )
    parser.add_argument(
        "-i",
        "--inputDir",
        help="Path to the folder where the input .xml files are stored",
        type=str,
    )
    parser.add_argument(
        "-o", "--outputFile", help="Name of output .csv file (including path)", type=str
    )

    parser.add_argument(
        "-l",
        "--labelMapDir",
        help="Directory path to save label_map.pbtxt file is specified.",
        type=str,
        default="",
    )

    args = parser.parse_args()

    if args.inputDir is None:
        args.inputDir = os.getcwd()
    if args.outputFile is None:
        args.outputFile = args.inputDir + "/labels.csv"

    assert os.path.isdir(args.inputDir)
    os.makedirs(os.path.dirname(args.outputFile), exist_ok=True)
    xml_df, classes_names = xml_to_csv(args.inputDir)
    xml_df.to_csv(args.outputFile, index=None)
    print("Successfully converted xml to csv.")
    if args.labelMapDir:
        os.makedirs(args.labelMapDir, exist_ok=True)
        label_map_path = os.path.join(args.labelMapDir, "label_map.pbtxt")
        print("Generate `{}`".format(label_map_path))

        # Create the `label_map.pbtxt` file
        pbtxt_content = ""
        for i, class_name in enumerate(classes_names):
            pbtxt_content = (
                pbtxt_content
                + "item {{\n    id: {0}\n    name: '{1}'\n}}\n\n".format(
                    i + 1, class_name
                )
            )
        pbtxt_content = pbtxt_content.strip()
        with open(label_map_path, "w") as f:
            f.write(pbtxt_content)


if __name__ == "__main__":
    main()
