# Description: Filter out the bounding boxes based on the following criteria:
# 1. Huge bounding boxes: Filter out the bounding boxes whose area is larger than a certain threshold.
# 2. Non-maximum suppression (NMS): Filter out the bounding boxes that have a high IoU with other bounding boxes.
# 3. Category: Filter out the bounding boxes whose category is in the FILTER_OUT_CATEGORY list.
#
# Usage: python filter_box.py -i <input_file> -o <output_file> [-n] [-hu] [-c]
#   -i, --input: Path to the input JSON file.
#   -o, --output: Path to the output JSON file.
#   -n, --nms: Perform non-maximum suppression.
#   -hu, --huge: Filter out huge bounding boxes.
#   -c, --category: Filter out bounding boxes based on category.

import json
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, required=True)
parser.add_argument('-o', '--output', type=str, required=True)
parser.add_argument('-n', '--nms', action='store_true')
parser.add_argument('-hu', '--huge', action='store_true')
parser.add_argument('-c', '--category', action='store_true')

IMG_SIZE = 960*540
HUGE_THRESHOLD = 0.9
NMS_THRESHOLD = 0.7
FILTER_OUT_CATEGORY = ['controller', 'controllers', 'hand', 'hands']

def IoU(box1, box2):
    x1, y1, w1, h1 = box1
    x2, y2, w2, h2 = box2

    inter_x1 = max(x1, x2)
    inter_y1 = max(y1, y2)
    inter_x2 = min(x1 + w1, x2 + w2)
    inter_y2 = min(y1 + h1, y2 + h2)
    
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
    union_area = w1 * h1 + w2 * h2 - inter_area
    iou = inter_area / union_area if union_area > 0 else 0
    
    return iou


def filter_huge(preds):
    filtered_preds = []
    for pred in preds:
        x, y, w, h = pred['bbox']
        if w * h / IMG_SIZE < HUGE_THRESHOLD:
            filtered_preds.append(pred)
    return filtered_preds


def nms(preds):
    sorted_preds = sorted(preds, key=lambda x: x['score'], reverse=True)
    nms_preds = []
    for pred in sorted_preds:
        if len(nms_preds) == 0:
            nms_preds.append(pred)
            continue
        for nms_pred in nms_preds:
            if IoU(pred['bbox'], nms_pred['bbox']) > NMS_THRESHOLD:
                break
        else:
            nms_preds.append(pred)
    return nms_preds


def filter_category(preds):
    filtered_preds = []
    for pred in preds:
        if pred['category_id'].lower() not in FILTER_OUT_CATEGORY:
            filtered_preds.append(pred)
    return filtered_preds


def main(args):
    with open(args.input, 'r') as f:
        pred_data = json.load(f)

    img_pred = {}

    for pred in pred_data:
        img_id = pred['image_id']
        if img_id not in img_pred:
            img_pred[img_id] = []
        img_pred[img_id].append(pred)

    pred_data = []
    
    for img_id in img_pred:
        preds = img_pred[img_id]
        if args.huge:
            preds = filter_huge(preds)
        if args.nms:
            preds = nms(preds)
        if args.category:
            preds = filter_category(preds)

        pred_data.extend(preds)

    with open(args.output, 'w') as f:
        json.dump(pred_data, f)


if __name__ == '__main__':
    args = parser.parse_args()
    main(args)