#!/usr/bin/env python

import os
import lxml.etree as ET
import ast
import help_parser as hp

CLASSES = ['car', 'person', 'bicycle']
KITTI_C = ['Car', 'Pedestrian', 'Cyclist']

def parse_args():
	parser = hp.HelpParser(description='Helper program to facilitate XML format and KITTI format')
	parser.add_argument('xml', help='Input XML file from recognize tool')
	parser.add_argument('dir', help='Output folder to store KITTI files')
	parser.add_argument('-v', default=False, help='Verbose mode', action='store_true')

	args = parser.parse_args()
	return args


def convert_data(data, args):
	key = 'path'
	file = os.path.splitext(os.path.basename(data.get(key)))[0] + '.txt'
	if args.v:
		print file
	with open(os.path.join(args.dir, file), mode='w') as f:
		for object in data.iter('object'):
			if object.get('class') not in CLASSES:
				continue
			conf = object.get('prob')
			cl = KITTI_C[CLASSES.index(object.get('class'))]
			bbox = ast.literal_eval(object.get('bbox'))
			s = cl + ' 0.0'*3 + ' ' + ' '.join(map(str,bbox)) + ' 0.0'*7 + ' ' + str(conf) + '\n'
			f.write(s)

if __name__ == '__main__':
	args = parse_args()
	xml = ET.parse(args.xml)
	root = xml.getroot()
	if root.get('type') != 'folder':
		raise ValueError('You can process only XML file coming from \'folder\' type experiment!')
	
	if not os.path.isdir(args.dir):
		os.makedirs(args.dir)
	
	for child in root.iter('file'):
		convert_data(child, args)
