"""
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, glob, sys, json
import pandas as pd
import argparse
import xml.etree.ElementTree as ET

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, f'{current_dir}/../')
from quannh.yolo_func import get_labels
labels = get_labels('insects')

def xml_to_csv(path,percentage_test=None):
    """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
    """
    counter = 1
    index_test = round(100 / percentage_test)
    is_test = False

    xml_list = []
    xml_test = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            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)
                    )
            xml_list.append(value)

    for json_file in glob.glob(path + '/*.json'):
        with open(json_file, 'r') as f:
            datastore = json.load(f)
            width = datastore['imageWidth']
            if counter == index_test:
                counter = 1
                is_test = True
            else:
                is_test = False
                counter = counter + 1

            for shape in datastore["shapes"]:
                if shape['label'] in labels:
                    value = (datastore["imagePath"],
                            int(datastore['imageWidth']),
                            int(datastore['imageHeight']),
                            shape['label'],
                            int(shape['points'][0][0]),
                            int(shape['points'][0][1]),
                            int(shape['points'][1][0]),
                            int(shape['points'][1][1])
                            )
                    if is_test:
                        xml_test.append(value)
                    else:
                        xml_list.append(value)

    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    if percentage_test is None:
        return pd.DataFrame(xml_list, columns=column_name)
    else :
        return pd.DataFrame(xml_list, columns=column_name), pd.DataFrame(xml_test, columns=column_name)
        


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("-d",
                        "--outputDirectory",
                        help="Name of output directory of .csv file (including path)", type=str)
    args = parser.parse_args()

    if(args.inputDir is None):
        args.inputDir = os.getcwd()

    assert(os.path.isdir(args.inputDir))

    if(args.outputFile is None):
        args.outputFile = args.inputDir + "/labels.csv"
    else :
        xml_df = xml_to_csv(args.inputDir)
        xml_df.to_csv(args.outputFile, index=None)
    
    if args.outputDirectory is not None:
        xml_train, xml_test = xml_to_csv(args.inputDir,20)
        xml_train.to_csv(args.outputDirectory + "/train_labels.csv", index=None)
        xml_test.to_csv(args.outputDirectory + "/test_labels.csv", index=None)
        

    
    
    # xml_df.to_csv(args.outputFile, index=None)
    print('Successfully converted xml to csv.')


if __name__ == '__main__':
    main()