import argparse
import os
from shapely.geometry import Polygon

def markup_load(path):
    markup_file = open(path, 'r')

    markup = []
    for line in markup_file:
        line = line.strip().split()
        ptype = line[0]
        if ptype != 'bld':
            continue
        count = int(line[1])
        points = []
        for i in range(count):
            x = float(line[2+2*i])
            y = float(line[2+2*i+1])
            points.append((x, y))
        markup.append(Polygon(points).buffer(0))
    return markup


def iou(polygon1, polygon2):
    intersection = polygon1.intersection(polygon2)
    if intersection.area > 0:
        union = polygon1.union(polygon2)
        return intersection.area/union.area
    else:
        return 0


def print_result(gt_cnt, test_cnt, tp_cnt):
    not_fnd_cnt = gt_cnt - tp_cnt
    false_fnd_cnt = test_cnt - tp_cnt
    precision  = tp_cnt/float(test_cnt) if test_cnt > 0 else 1.0
    recall = tp_cnt/float(gt_cnt) if gt_cnt > 0 else 1.0

    print '| bld  |{:^7d}|{:^7d}|{:^11d}|{:^13d}|{:^11f}|{:^8f}|'.format(gt_cnt, tp_cnt, not_fnd_cnt, false_fnd_cnt, precision, recall)


def compare(gt_file, test_file, iou_threshold):
    gt_markup = markup_load(gt_file)
    test_markup = markup_load(test_file)

    gt_cnt = len(gt_markup)
    test_cnt = len(test_markup)
    iou_arr = []
    for gt_idx in range(gt_cnt):
        for test_idx in range(test_cnt):
            iou_value = iou(gt_markup[gt_idx], test_markup[test_idx])
            if iou_value > iou_threshold:
                iou_arr.append([iou_value, gt_idx, test_idx])

    iou_arr.sort(key=lambda x: x[0], reverse=True)

    test_fnd = [False for polygon in test_markup]
    gt_fnd   = [False for polygon in gt_markup]

    for i in range(len(iou_arr)):
        gt_idx = iou_arr[i][1]
        if (gt_fnd[gt_idx]):
            continue
        test_idx = iou_arr[i][2]
        if (test_fnd[test_idx]):
            continue
        gt_fnd[gt_idx] = True
        test_fnd[test_idx] = True

    return gt_cnt, test_cnt, sum(gt_fnd)


def compare_file(gt_file, test_file, iou_threshold):
    gt_cnt, test_cnt, tp_cnt = compare(gt_file, test_file, iou_threshold)
    print '| type | count | found | not_found | false_found | precision | recall |'
    print_result(gt_cnt, test_cnt, tp_cnt)


def compare_list(gt_path, test_path, iou_threshold):
    gt_filenames = os.listdir(gt_path)

    print '| type | count | found | not_found | false_found | precision | recall |'
    total_tp_cnt = 0
    total_gt_cnt = 0
    total_test_cnt = 0
    for filename in sorted(gt_filenames):
        gt_file = os.path.join(gt_path, filename)
        test_file = os.path.join(test_path, filename)
        if not os.path.isfile(test_file):
            print 'There is no file to compare with {}'.format(filename)
            continue

        gt_cnt, test_cnt, tp_cnt = compare(gt_file, test_file, iou_threshold)
        print filename
        print_result(gt_cnt, test_cnt, tp_cnt)
        total_gt_cnt += gt_cnt
        total_test_cnt += test_cnt
        total_tp_cnt += tp_cnt

    print 'total'
    print_result(total_gt_cnt, total_test_cnt, total_tp_cnt)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--gt', required = True, help='Path to test file/folder with polygons')
    parser.add_argument('--test', required = True, help='Path to ground truth file/folder with polygons')
    parser.add_argument('--iou_threshold', type = float, default = 0.5, help = 'Thresh value for IOU')
    args = parser.parse_args()

    if os.path.isfile(args.gt):
        compare_file(args.gt, args.test, args.iou_threshold)
    else:
        compare_list(args.gt, args.test, args.iou_threshold)


if __name__ == '__main__':
    main()
