#!/usr/bin/env python

import os
import lxml.etree as ET
import ast
import help_parser as hp
import warnings
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

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='Create Precision-recall curve')
	parser.add_argument('out', help='Name of output file for plotting curves, if no extension is given, file will be saved as eps')
	parser.add_argument('-gt', nargs='+', help='Name of input files with ground truth')
	parser.add_argument('-pred', nargs='+', help='Name of input files with predictions')
	parser.add_argument('-bins', dest='num_bins', help='Number of bins for recall values', default=41, type=int)
	parser.add_argument('-over', help='Overlap ratio for a prediction to be considered correct [0.5]', default=0.5, type=float)
	parser.add_argument('-data', help='Text file to sava plotting data for future use in different plots', default=None)
	parser.add_argument('-cls', dest='classes', help='Compute statistics for distinct classes', default=False, action='store_true')
	parser.add_argument('-v', dest='verbose', help='Verbose mode', default=False, action='store_true')

	args = parser.parse_args()
	return args


def parse_xml(input, pred=False):
	xml = ET.Element('pred' if pred else 'gt')
	for f in input:
		root = ET.parse(f).getroot()[0]
		if root.tag != 'folder':
			warnings.warn('File ' + f + ' is not an experiment run on folder, skipping!')
			continue
		xml.extend(root)
	return xml


def get_threshs(scores, num_of_pred, bins):
	scores.sort(reverse=True)
	cur = 0
	ret = list()
	for i, _ in enumerate(scores):
		l = (i+1.0)/num_of_pred
		if i == len(scores) - 1:
			r = l
		else:
			r = (i+2.0)/num_of_pred
		if (r - cur) < (cur - l) and i != len(scores) - 1:
			continue
		ret.append(i)
		cur += 1.0/(bins-1)
	return scores, ret


def iou(bbox1, bbox2):
	x11, y11, x12, y12 = bbox1
	x21, y21, x22, y22 = bbox2
	xo1 = max(x11, x21)
	xo2 = min(x12, x22)
	yo1 = max(y11, y21)
	yo2 = min(y12, y22)
	
	wo = xo2 - xo1
	ho = yo2 - yo1
	
	if wo <= 0 or ho <= 0:
		return 0
	
	inter = wo*ho
	a = (x12-x11)*(y12-y11)
	b = (x22-x21)*(y22-y21)
	return inter/(a+b-inter)
	

def shorten_by_class(input, cls, split):
	ret = dict()
	if split:
		for c in cls:
			ret[c] = list()
	else:
		ret['all'] = list()
	for gt, pred in input:
		if split:
			for c in cls:
				gc, _ = split_by_func(gt, lambda x: x['class'] == c)
				pc, _ = split_by_func(pred, lambda x: x['class'] == c)
				ret[c].append((gc, pc))
		else:
			pp, _ = split_by_func(pred, lambda x: x['class'] in cls)
			ret['all'].append((gt, pp))
	return ret


def pairings(gt, pred, args):
	ret = list()
	scores = list()
	cls=list()
	num_of_pred = 0
	lc = len(gt)
	for i, child in enumerate(gt):
		name = os.path.basename(child.get('path'))
		cls += [o.get('class') for o in child.xpath('./object')]
		if args.verbose:
			print str(i + 1) + '/' + str(lc) + ': pairing and preprocessing of ' + name
		path = './file[contains(@path, \''+name+'\')]'
		match = pred.xpath(path)
		if len(match) == 0:
			warnings.warn('There is no prediction for groundtruth of file '+name+'. Skipping this file, won\'t count into statistics')
			continue
		cc = list()
		for c in child:
			d = dict(c.attrib)
			d['bbox'] = ast.literal_eval(d['bbox'])
			cc.append(d)
		mm = list()
		if len(match) >= 1:
			for m in match:
				mmm = list()
				num_of_pred += len(m)				
				for dd in m:
					d = dict(dd.attrib)
					d['bbox'] = ast.literal_eval(d['bbox'])
					scores.append(float(d['prob']))
					d['prob'] = float(d['prob'])
					mmm.append(d)
				mm.extend(mmm)
		ret.append((cc, mm))
	
	cls = set(cls)
	
	ret = shorten_by_class(ret, cls, args.classes)

	return ret, get_threshs(scores, num_of_pred, args.num_bins)


def split_by_func(lst, func):
	acc, rej = [], []
	for x in lst:
		if func(x):
			acc.append(x)
		else:
			rej.append(x)
	return acc, rej


def compute_stats(pairs, bins, args):
	files = dict()
	for c in pairs.keys():
		if args.verbose:
			print 'Computing stats for class ' + c
		files[c] = list()
		for (gt, pred) in pairs[c]:
			matches = list()
			for i, g in enumerate(gt):
				found = False
				g_cl = g['class']
				g_bbox = g['bbox']
				for p in pred:
					if p['class'] != g_cl:
						continue
					if iou(g_bbox, p['bbox']) >= args.over:
						found = True
						pred.remove(p)
						matches.append(p['prob'])
						break
				if not found:
					matches.append(-1)
			rest = [p['prob'] for p in pred]
			files[c].append((matches, rest))
		if args.verbose:
			print 'Stats for class ' + c + ' computed'
	ret = dict()
	pr = dict()
	scores, bb = bins
	for c in pairs.keys():
		ret[c] = list()
		pr[c] = list()
		for i, b in enumerate(bb):
			if args.verbose:
				print 'Computing stats on ' + str(i + 1) +'th bin (' + str(scores[b]) + ') of class ' + c
			stat = get_stats(files[c], scores[b])
			nb = b + int(len(scores)/200.0)
			while (stat['tp'] == 0 or stat['fp'] == 0) and nb < len(scores) - 1:
				if args.verbose:
					print 'No stats for class ' + c + ' and threshold ' + str(scores[b]) + ', trying threshold ' + str(scores[nb])
				b = nb
				stat = get_stats(files[c], scores[b])
				nb = b + int(len(scores)/200.0)
			if not (stat['tp'] == 0 or stat['fp'] == 0):
				stat['pr'] = float(stat['tp'])/(stat['tp'] + stat['fp'])
				stat['re'] = float(stat['tp'])/(stat['tp'] + stat['fn'])
				ret[c].append(stat)
				pr[c].append(stat['pr'])
			else:
				break

		for i in range(len(ret[c])):
			ret[c][i]['pr'] = max(pr[c][i:])
	return ret


def get_stats(files, b):
	stat = dict()
	stat['tp'] = 0
	stat['fn'] = 0
	stat['fp'] = 0
	stat['tr'] = b
	for matches, rest in files:
		l, m = split_by_func(matches, lambda x: x < b)
		stat['tp'] += len(m)
		stat['fn'] += len(l)
		_, m = split_by_func(rest, lambda x: x < b)
		stat['fp'] += len(m)
	return stat


def plot(stats, args):
	
	fig = plt.figure()
	pplt = fig.add_subplot(111)
	
	if args.data is not None:
		f = open(args.data, mode='w')
	
	for c in stats.keys():
		
		pr = [s['pr'] for s in stats[c]]
		re = [s['re'] for s in stats[c]]
		tr = [s['tr'] for s in stats[c]]
		pplt.plot(re, pr, label=c)
		if args.data is not None:
			f.write(c + '\n')
			f.write(' '.join(map(str,re)) + '\n')
			f.write(' '.join(map(str,pr)) + '\n')
	
	if args.data is not None:
		f.close()
	pplt.set_xlabel('Recall')
	pplt.set_ylabel('Precision')
	pplt.axis([0,1,0,1])
	title = 'Precision-recall curve'
	for pr in args.pred:
		title += '\n'+shorten_string('Pred: '+os.path.abspath(pr), 36)
	pplt.set_title(title)
	pplt.legend()
	plt.tight_layout()
	fig.savefig(args.out)
	plt.close(fig)


def shorten_string(string, max_len):
	if len(string) <= max_len:
		return string
	phase = (max_len-2)/2
	return string[0:phase] + '..' + string[-phase:]
	

if __name__ == '__main__':
	args = parse_args()
	
	bname = os.path.basename(args.out)
	if bname == os.path.splitext(bname)[0]:	
		args.out = args.out + '.eps'
	
	if os.path.isfile(args.out):
		warnings.warn('File ' + args.out + ' will be overwritten!')

	gt = parse_xml(args.gt)
	if args.verbose:
		print 'Ground truth files loaded'
		print str(len(gt)) + ' image files'
	pred = parse_xml(args.pred, True)
	if args.verbose:
		print 'Prediction files loaded'
		print str(len(pred)) + ' image files'
	
	pairs, bins = pairings(gt, pred, args)
	if args.verbose:
		print 'Paired files and preprocessing done'
	stats = compute_stats(pairs, bins, args)
	if args.verbose:
		print 'Stats computed'
	plot(stats, args)
