#!/usr/bin/env python

import _init_paths
from fast_rcnn.config import cfg
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
from utils.timer import Timer
import cv2
import caffe
import help_parser as hp
import os
import numpy as np
import warnings
import lxml.etree as ET
import ast
from xml.dom import minidom

MODEL_DIR = _init_paths.MODEL_DIR

CLASSES = ('__bg__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair',
	'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa','train', 'tvmonitor')

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='Recognize objects in video')
	parser.add_argument('input', help='Name of either a video or a folder with images to process')
	parser.add_argument('-folder', dest='folder', default=False, help='Folder to save output data [folder of input data]')
	parser.add_argument('-prefix', dest='prefix', default=False, help='Prefix of new files, required when input is folder and you haven\'t specified output folder')
	parser.add_argument('-ox', dest='output_xml', help='XML output file [stdout]', default=False)
	parser.add_argument('-t', dest='timing', default=False, action='store_true', help='Include timing data in outputted XML file [False]')
	parser.add_argument('-gpu',type=int, dest='gpu_ID', default='0', help='GPU device ID to be used [0]')
	parser.add_argument('-cpu', dest='cpu', action='store_true', help='Use CPU only mode (overrides -gpu)')
	inside = sorted(os.listdir(MODEL_DIR))
	nets = [f for f in inside if os.path.isdir(os.path.join(MODEL_DIR,f)) and f != 'models']
	parser.add_argument('-net', dest='net', help='Network to use [vgg16_orig]', choices=nets, default='vgg16_orig')
	parser.add_argument('-conf', dest='conf', help='Confidence threshold [0.8]', default=0.8, type=float)
	parser.add_argument('-move', dest='move', help='Threshold of movement of objects between frames as a percentage [0.02]. Used only for videos.', default=0.02, type=float)
	parser.add_argument('-nms', dest='nms', help='NMS threshold [0.3]', default=0.3, type=float)
	parser.add_argument('-box', dest='box', action='store_false', help='Show percentages and classes in bounding boxes [True]', default=True)
	parser.add_argument('-nv', dest='visual', action='store_false', default=True, help='No visualisation output (videos and images)')	
	parser.add_argument('-vaf', dest='vaf', action='store_true', default=False, help='Treat videos as folders [False]')
	parser.add_argument('-v', dest='verbose', action='store_true', default=False, help='Verbose mode')
	args = parser.parse_args()

	if args.visual:
		if os.path.isdir(args.input) and (not args.folder or args.folder == args.input) and not args.prefix:
			raise ValueError('When input is directory, you have to specify either output folder or prefix!')
	
		if os.path.isfile(args.input) and not args.prefix:
			raise ValueError('When input is video file, you have to specify prefix for newly saved video!')

	return args

def visualise_frame(im, frame, show_info=False):
	for item in frame:
		cv2.rectangle(im, (int(item[2][0]), int(item[2][1])), (int(item[2][2]), int(item[2][3])), (0,0,255), 3)
		if show_info:
			text = item[0] + ': ' + '%.3f' % item[1]
			size, _ = cv2.getTextSize(text, cv2.cv.CV_FONT_HERSHEY_SIMPLEX, 0.6, 1)
			center = (int((item[2][2] + item[2][0])/2), int(item[2][1]))
			cv2.rectangle(im, (int(center[0] - size[0]/2), center[1] - size[1] if center[1] >= size[1] else 0), (int(center[0] + size[0]/2), center[1] if center[1] >= size[1] else size[1]), (0,0,255), -1)
			cv2.putText(im, text, (int(center[0] - size[0]/2), center[1]  if center[1] >= size[1] else center[1] + size[1]), cv2.cv.CV_FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), bottomLeftOrigin=False)
	return im

def process_image(net, im, args, video=True):
	if args.timing:
		timer = Timer()
		timer.tic()
	scores, boxes = im_detect(net, im)
	frame = list()
	for i in range(1,len(CLASSES)):
		i_boxes = boxes[:,4*i:4*(i+1)]
		i_scores = scores[:,i]
		dets = np.hstack((i_boxes, i_scores[:, np.newaxis])).astype(np.float32)
		dets = dets[nms(dets, args.nms),:]
		inds = np.where(dets[:,-1] >= args.conf)[0]
		if len(inds) == 0:
			continue
		for j in inds:
			frame.append((CLASSES[i], dets[j, -1], (dets[j,0], dets[j,1], dets[j,2], dets[j,3])))
	if args.timing:
		timer.toc()
	if args.visual:
		new_im = visualise_frame(im, frame, args.box)
	if not video:
		elem = ET.Element('file')
		if args.timing:
			elem.set('time', str(timer.total_time))
		for f in frame:
			a = ET.Element('object')
			a.set('class', str(f[0]))
			a.set('prob', str(f[1]))
			a.set('bbox', str(f[2]))
			elem.append(a)
	return (frame, new_im if args.visual else None, elem if not video else None, timer.total_time if args.timing else 0)


def run_folder(net, folder, args):
	files = [f for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))]
	files.sort()
	count = len(files)
	xml = ET.Element('folder')
	xml.set('path', folder)
	xml.set('files', str(count))
	x = 0
	for i, f in enumerate(files):
		im = cv2.imread(os.path.join(folder, f))
		if im is not None:
			if args.verbose:
				print(str(i+1)+'/'+str(count)+': '+os.path.join(folder, f))
			frame, vis_im, elem, _ = process_image(net, im, args, False)
			elem.set('path', os.path.join(folder, f))
			if args.visual:
				elem.set('newpath', os.path.join(args.folder, args.prefix + f))
			xml.append(elem)
			if args.visual:
				cv2.imwrite(os.path.join(args.folder, args.prefix + f), vis_im)
		else:
			try:
				video = cv2.VideoCapture(os.path.join(folder, f))
				if args.verbose:
					print(str(i+1) + '/' + str(count) + ': ' + os.path.join(folder, f) + ' is video file')
					elem = run_video(net, video, args, f)
					xml.append(elem)
			except cv2.error:
				x += 1
				if args.verbose:
					print(str(i+1) + '/' + str(count) + ': ' + os.path.join(folder, f) + ' is not an image or video file, skipping')
	xml.set('usable_files', str(count - x))
	return xml
	

def run_video(net, video, args, filename):
	objects = list()
	count = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
	i = 1
	size = (int(video.get(3)), int(video.get(4)))
	if args.visual:
		n_video = cv2.VideoWriter(os.path.join(args.folder, args.prefix + filename), int(video.get(6)), int(video.get(5)), size)
	time = 0
	while video.isOpened():
		ret, im = video.read()
		if not ret:
			break
		if args.verbose:
			print filename + ': ' + str(i) + '/' + str(count) + ' (number of frames might be incorrect)'
		frame, vis_im, _, t = process_image(net, im, args)
		time += t
		if args.visual:
			n_video.write(vis_im)
		objects.append(frame)
		i = i + 1
	video.release()
	xml = create_video_xml(objects, args, filename, size)
	if args.visual:
		n_video.release()
		xml.set('newpath', os.path.join(args.folder, args.prefix + filename))
	if args.timing:
		xml.set('avg_time', str((time/(i-1))))
	return xml


def create_video_xml(objects, args, filename, size):
	size = size*2
	size = tuple([a * args.move for a in size])
	xml = ET.Element('file')
	xml.set('path', os.path.abspath(args.input) if os.path.isfile(args.input) else os.path.join(args.input, filename))
	xml.set('framecount', str(len(objects)))
	if args.vaf:
		for i, frame in enumerate(objects):
			for obj in frame:
					elem = ET.Element('object')
					elem.set('class', obj[0])
					elem.set('frame',str(i))
					elem.set('prob', str(obj[0]))
					elem.set('bbox', str(obj[2]))
					xml.append(elem)
	else:
		for i, frame in enumerate(objects):
			last = xml.xpath('./object[@last_frame='+str(i-1)+']')
			for obj in frame:
				found = None
				for elem in last:
					if obj[0] != elem.get('class'):
						continue
					elem_bb = ast.literal_eval(elem.get('last_bbox'))
					diff = tuple([abs(a - b) < c for a, b, c in zip(elem_bb, obj[2], size)])
					if all(diff):
						found = elem
						break
					else:
						continue
				
				if found is not None:
					found.set('last_bbox', str(obj[2]))
					found.set('last_frame', str(i))
					high = float(found.get('highest'))
					if obj[1] > high:
						found.set('highest', str(obj[1]))
					else:
						low = float(found.get('lowest'))
						if obj[1] < low:
							found.set('lowest', str(obj[1]))
				else:
					elem = ET.Element('object')
					elem.set('class', obj[0])
					elem.set('first_frame', str(i))
					elem.set('last_frame', str(i))
					elem.set('first_bbox', str(obj[2]))
					elem.set('last_bbox', str(obj[2]))
					elem.set('highest', str(obj[1]))
					elem.set('lowest', str(obj[1]))
					xml.append(elem)
	return xml
				

if __name__ == '__main__':
	args = parse_args()
	
	if os.path.isdir(args.input):
		args.input = os.path.abspath(args.input)
		if not args.folder:
			args.folder = args.input
	else:
		if not args.folder:
			args.folder = os.path.dirname(os.path.abspath(args.input))

	if args.output_xml and os.path.isfile(args.output_xml):
		warnings.warn('File '+args.output_xml+' will be overwritten!')
	
	if args.output_xml and not os.path.exists(os.path.dirname(args.output_xml)):
		os.makedirs(os.path.dirname(args.output_xml))
	
	if not os.path.isdir(args.folder):
		os.makedirs(args.folder)
	
	if not args.prefix:
		args.prefix=''

	cfg.TEST.HAS_RPN = True
	if args.cpu:
		caffe.set_mode_cpu()
	else:
		caffe.set_mode_gpu()
		caffe.set_device(args.gpu_ID)
		cfg.GPU_ID = args.gpu_ID
	
	proto = os.path.join(MODEL_DIR, args.net, 'test.pt')
	model = os.path.join(MODEL_DIR, args.net, args.net + '.caffemodel')
	if not os.path.isfile(proto) or not os.path.isfile(model):
		raise IOError('Problem with finding net files!')
	
	timer = Timer()
	timer.tic()
	net = caffe.Net(proto, model, caffe.TEST)
	if os.path.isdir(args.input):
		xml = run_folder(net,args.input, args)
	else:
		try:
			video = cv2.VideoCapture(args.input)
			xml = run_video(net, video, args, os.path.basename(args.input))
		except cv2.error:
			raise TypeError('File '+args.input+' is not a valid video file!')
	timer.toc()
	if xml is not None:
		root = ET.Element('Faster_RCNN_experiment')
		root.set('type', 'folder' if os.path.isdir(args.input) else 'video')
		root.set('network', args.net)
		root.set('nms', str(args.nms))
		root.set('conf', str(args.conf))
		root.set('move', str(args.move))
		root.set('time', str(timer.total_time))
		root.append(xml)
		rough = ET.tostring(root)
		if args.output_xml:
			with open(args.output_xml, mode='w') as f:
				minidom.parseString(rough).writexml(f, addindent='\t', newl='\n')
		else:				
			minidom.parseString(rough).writexml(sys.stdout, addindent='\t', newl='\n')
