import json
import random
from tqdm import tqdm
import argparse
import numpy as np

MAX_INTERACT_PER_IMAGE = 60
WIDTH = 960
HEIGHT = 540


def in_pred_bbox(img_preds, X, Y):
    for pred in img_preds:
        if X - pred['bbox'][0] >= 0 and \
            Y - pred['bbox'][1] >= 0 and \
            X - pred['bbox'][0] <= pred['bbox'][2] and \
            Y - pred['bbox'][1] <= pred['bbox'][3]:
            return True
    return False


def gen_interact_random(annotation_data, width, height):
    random_inter = {}
    # Batch-by-batch generation, 1 point per image in each batch
    # for k in tqdm(range(1, MAX_INTERACT_PER_IMAGE + 1), desc='Randomly generating interact points'):
    for k in range(1, MAX_INTERACT_PER_IMAGE + 1):
        cur_group = []
        for img in annotation_data['images']:
            img_id = img['id']
            cur_group.append({
                'X': random.randint(1, width),
                'Y': random.randint(1, height),
                'img_id': img_id
            })
        random_inter[k] = cur_group
    return random_inter


def gen_interact_guided(anno_data, guidance, width, height):
    interact = {}
    img_pred = {}
    for pred in guidance:
        img_id = pred['image_id']
        if img_id not in img_pred:
            img_pred[img_id] = []
        img_pred[img_id].append(pred)
    img_pred_mask = {}
    for img_id in img_pred.keys():
        img_pred_mask[img_id] = np.zeros((width, height))
        if not img_pred[img_id]:
            raise ValueError(f'No prediction for image id {img_id} in guidance')
        for pred in img_pred[img_id]:
            x, y, w, h = pred['bbox']
            x = np.floor(x).astype(int)
            y = np.floor(y).astype(int)
            w = np.ceil(w).astype(int)
            h = np.ceil(h).astype(int)
            img_pred_mask[img_id][x:x+w, y:y+h] = 1
    # Batch-by-batch generation, 1 point per image in each batch
    # for k in tqdm(range(1, MAX_INTERACT_PER_IMAGE + 1), desc='Generating guided interact points'):
    for k in range(1, MAX_INTERACT_PER_IMAGE + 1):
        outside_prob = k / MAX_INTERACT_PER_IMAGE
        cur_group = []
        for img in anno_data['images']:
            img_id = img['id']
            outside = random.random() < outside_prob
            # if outside or img_id not in img_pred.keys() or np.sum(img_pred_mask[img_id]) == 0:
            #     X = random.randint(1, width)
            #     Y = random.randint(1, height)
            # else:
            #     Xs, Ys = np.where(img_pred_mask[img_id] == 1)
            #     idx = random.randint(0, len(Xs) - 1)
            #     X = int(Xs[idx])
            #     Y = int(Ys[idx])
            for _ in range(10000): # To avoid infinite loop
                X = random.randint(1, width)
                Y = random.randint(1, height)
                if outside or img_id not in img_pred.keys() or in_pred_bbox(img_pred[img_id], X, Y):
                    break
            cur_group.append({
                'X': X,
                'Y': Y,
                'img_id': img_id
            })
        interact[k] = cur_group
    return interact


def main(args):
    with open(args.annotation, 'r') as f:
        anno_data = json.load(f)
    if args.guidance:
        with open(args.guidance, 'r') as f:
            guidance = json.load(f)
        guidance_img = [pred['image_id'] for pred in guidance]
        anno_data['images'] = [img for img in anno_data['images'] if img['id'] in guidance_img]
        anno_data['annotations'] = [anno for anno in anno_data['annotations'] if anno['image_id'] in guidance_img]
        interact_data = gen_interact_guided(anno_data, guidance, WIDTH, HEIGHT)
    else:
        interact_data = gen_interact_random(anno_data, WIDTH, HEIGHT)
    with open(args.output, 'w') as f:
        json.dump(interact_data, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate interact points')
    parser.add_argument('-a', '--annotation', type=str, required=True, help='Annotation file path')
    parser.add_argument('-g', '--guidance', type=str, help='guidance file path')
    parser.add_argument('-o', '--output', type=str, required=True, help='Output file path')
    args = parser.parse_args()
    main(args)