#!/usr/bin/env python

import os
import lxml.etree as ET
import ast
import help_parser as hp
from xml.dom import minidom
import warnings

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

fw_orig = warnings.formatwarning
warnings.formatwarning = lambda msg, categ, fname, lineno, line=None: fw_orig(msg, categ, fname, lineno, '')

def parse_args():
	parser = hp.HelpParser(description='Helper program to facilitate XML format and KITTI format')
	parser.add_argument('folder', help='Input folder of KITTI labels')
	parser.add_argument('xml', help='Output XML file formatted as from recognize tool')
	parser.add_argument('-im_dir', default=None, help='Image folder of KITTI files used in XML, usually it differs from label folder. [label folder]')
	parser.add_argument('-im_ext', default='jpg', help='Extension of images, input without a dot [jpg]')
	parser.add_argument('-v', default=False, help='Verbose mode', action='store_true')

	args = parser.parse_args()
	
	args.folder = os.path.abspath(args.folder)
	
	if args.im_dir is None:
		args.im_dir = args.folder
	return args

def convert_data(fname, args):
	xml = ET.Element('file')
	xml.set('path', os.path.join(args.im_dir, os.path.splitext(os.path.basename(fname))[0] + '.' + args.im_ext))
	pred = False
	with open(fname, mode='r') as f:
		for line in f:
			elems = line.split(' ')
			if elems[0] not in KITTI_C:
				continue
			el = ET.Element('object')
			el.set('class', CLASSES[KITTI_C.index(elems[0])])
			if len(elems) == 16:
				el.set('prob', str(float(elems[15])))
				pred = True
			else:
				el.set('prob', 'gt')
			el.set('bbox', str(tuple(map(float,elems[4:8]))))
			xml.append(el)
	return xml, pred
		

if __name__ == '__main__':
	args = parse_args()
	if os.path.isfile(args.xml):
		warnings.warn('File ' + args.xml + ' will be overwritten')
	
	files = [f for f in os.listdir(args.folder) if os.path.isfile(os.path.join(args.folder, f))]
	files.sort()
	
	xml = ET.Element('folder')
	xml.set('path', args.im_dir)
	xml.set('files', str(len(files)))
	pred = False
	for f in files:
		if args.v:
			print f
		el, pr = convert_data(os.path.join(args.folder, f), args)
		pred = pred or pr
		if el is not None:
			xml.append(el)
	if pred:
		root = ET.Element('Faster_RCNN_experiment')
		root.set('type', 'folder')
		root.append(xml)
		xml = root
	else:
		root = ET.Element('Ground_truth')
		root.set('dataset', 'KITTI')
		root.append(xml)
		xml = root
	rough = ET.tostring(xml)
	with open(args.xml, mode='w') as f:
		minidom.parseString(rough).writexml(f, addindent='\t', newl='\n')