import os
import cv2
import json
import argparse
from tqdm import tqdm

ANNOTATION_FILE = '../../dataset/data/coco_merged/annotations/interactable.json'
# IMAGE_DIR = '../../dataset/data/coco_merged/images/interactable'
IMAGE_DIR = '../../dataset/data/visualize/result_ours_i'
INTERACT_FILE = './guided_interact.json'
OUTPUT_DIR = './output'

parser = argparse.ArgumentParser(description='Visualize interact points')
parser.add_argument('-a', '--annotation', type=str, default=ANNOTATION_FILE, help='Annotation file')
parser.add_argument('-i', '--image', type=str, default=IMAGE_DIR, help='Original image directory')
parser.add_argument('-int', '--interact', type=str, default=INTERACT_FILE, help='Interact points to visualize')
parser.add_argument('-o', '--output', type=str, default=OUTPUT_DIR, help='Output directory')


def main(args):
    os.makedirs(args.output, exist_ok=True)
    with open(args.annotation, 'r') as f:
        anno_data = json.load(f)
    with open(args.interact, 'r') as f:
        interact_data = json.load(f)
    for img in tqdm(anno_data['images'], desc='Visualizing interact points', total=len(anno_data['images'])):
        img_id = img['id']
        img_path = os.path.join('./result_ours_i', img['file_name'])
        if not os.path.exists(img_path):
            continue
        cimg = cv2.imread(img_path)
        for k in interact_data:
            for interact in interact_data[k]:
                if interact['img_id'] == img_id:
                    cv2.circle(cimg, (interact['X'], interact['Y']), 5, (0, 0, 100 + 125 // int(k)), -1)
        cv2.imwrite(os.path.join(args.output, img['file_name']), cimg)

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