import json
import os

control_pre = "This image contains "
control_post = "Based on this, "
    
def obj_ls2str(obj_ls):
    obj_str = ''
    for obj in obj_ls:
        obj_str += obj + ', '
    obj_str = obj_str[:-2] + '.'
    return obj_str


def coco_obj_ls(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None):
    assert prompt is not None
    if subset is not None:
        save_gt_obj_file = f'I{PROMPT_TYPE}_sub{len(subset)}_gt_obj'
    else:
        save_gt_obj_file = f'I{PROMPT_TYPE}_gt_obj'
    # if control is not None:
    #     save_gt_obj_file += '_control'
    save_gt_obj_file += '.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(path, coco_seg_file)
    coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]

    coco_gt_obj_ls = []
    for i in range(0, len(coco_data), 6):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)
        
        # if control is not None:
        #     obj_ls = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control + obj_ls2str(coco_seg_data[j]['objects'])}]}
        # else:
        obj_ls = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "objects": coco_seg_data[j]['objects']}
        coco_gt_obj_ls.append(obj_ls)

    os.makedirs(os.path.join(save_path, "obj"), exist_ok=True)
    
    save_gt_obj_path = os.path.join(save_path, "obj", save_gt_obj_file)
    with open(save_gt_obj_path, 'w') as f:
        json.dump(coco_gt_obj_ls, f)
        print(f"save at {save_gt_obj_path} successfully! (len: {len(coco_gt_obj_ls)})")
    

def coco_prompt_control(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None):
    assert prompt is not None
    if subset is not None:
        save_question_file = f'I{PROMPT_TYPE}_sub{len(subset)}'
    else:
        save_question_file = f'I{PROMPT_TYPE}'
    if control is not None:
        save_question_file += '_c2_new'
    save_question_file += '.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(path, coco_seg_file)
    coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]

    coco_questions_ls = []
    for i in range(0, len(coco_data), 6):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)
        
        if control is not None:
            # question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": control+obj_ls2str(coco_seg_data[j]['objects'])},{"from":"gpt", "value": "none"}]}
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control+obj_ls2str(coco_seg_data[j]['objects'])}]}
            # question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control[1] + obj_ls2str(coco_seg_data[j]['objects']).replace('.', ', ') + control[0]}]}
        else:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}]}
        coco_questions_ls.append(question)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))

    save_question_path = os.path.join(save_path, "question", save_question_file)
    with open(save_question_path, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {save_question_path} successfully! (len: {len(coco_questions_ls)})")
    

def coco_qa90_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None):
    assert prompt is not None
    detr_pred_file = "I4_pred_obj_detr_qa90_th0.9.json"
    question_file = "/data/linxi/workspace/LLaVA/playground/data/coco2014_val_gpt4_qa_30x3.jsonl"

    if subset is not None:
        save_question_file = f'I{PROMPT_TYPE}_sub{len(subset)}'
    else:
        save_question_file = f'qa{PROMPT_TYPE}'
    if control is not None:
        save_question_file += '_mmc4_grey' + '_'+ detr_pred_file.split('_')[-1].split('.json')[0]
        # save_question_file += '_c4_grey' + '_'+ detr_pred_file.split('_')[-1].split('.json')[0]
    save_question_file += '.json'

    # coco_data = [json.loads(q) for q in open(question_file, 'r')]
    with open (question_file, 'r') as f:
        # read lines
        coco_data = [json.loads(q) for q in f.readlines()]

    # coco_seg_path = os.path.join(path, coco_seg_file)
    # coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    # /data/linxi/workspace/POPE/llava_qa/obj/I4_pred_obj_detr.json
    coco_seg_path = os.path.join(save_path, "obj", detr_pred_file)
    coco_seg_data = json.load(open(coco_seg_path, 'r'))
    seg_image_ls = [seg['image'] for seg in coco_seg_data]
    print(len(coco_seg_data), coco_seg_data[0])
        
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]
# {"id": "000000525439", "image": "COCO_val2014_000000525439.jpg", "instruction": "What is the position of the skateboard in the image?", "output": "The skateboard in the image is in an upside-down position, with its wheels pointing up and laying on the ground.", "type": "conv"}

    coco_questions_ls = []
    for i in range(0, len(coco_data)):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        # while coco_data[i]['image'] != coco_seg_data[j]['image']:
        #     j = (j+1)%len(coco_seg_data)

        if coco_data[i]['image'] not in seg_image_ls:
            raise ValueError(f"image {coco_data[i]['image']} not found in seg data")
        j = seg_image_ls.index(coco_data[i]['image'])
        prompt = coco_data[i]['instruction']
        if control is not None:
            question = {"id": coco_data[i]['id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control[1] + obj_ls2str(coco_seg_data[j]['objects']).replace('.', ', ') + prompt.lower()}], "type": coco_data[i]['type']}
        else:
            question = {"id": coco_data[i]['id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}], "type": coco_data[i]['type']}
        coco_questions_ls.append(question)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))

    save_question_path = os.path.join(save_path, "question", save_question_file)
    with open(save_question_path, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {save_question_path} successfully! (len: {len(coco_questions_ls)})")


def coco_prompt_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None):
    assert prompt is not None
    detr_pred_file = "I4_pred_obj_detr_th0.95.json"

    if subset is not None:
        save_question_file = f'I{PROMPT_TYPE}_sub{len(subset)}'
    else:
        save_question_file = f'I{PROMPT_TYPE}'
    if control is not None:
        save_question_file += '_mmc4_grey' + '_'+ detr_pred_file.split('_')[-1].split('.json')[0]
        # save_question_file += '_c4_grey' + '_'+ detr_pred_file.split('_')[-1].split('.json')[0]
    save_question_file += '.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(path, coco_seg_file)
    coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    # /data/linxi/workspace/POPE/llava_qa/obj/I4_pred_obj_detr.json
    coco_seg_path = os.path.join(save_path, "obj", detr_pred_file)
    coco_seg_data = json.load(open(coco_seg_path, 'r'))
    print(len(coco_seg_data), coco_seg_data[0])
        
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]

    coco_questions_ls = []
    for i in range(0, len(coco_data), 6):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)
        
        if control is not None:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control[1] + obj_ls2str(coco_seg_data[j]['objects']).replace('.', ', ') + control[0]}]}
        else:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}]}
        coco_questions_ls.append(question)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))

    save_question_path = os.path.join(save_path, "question", save_question_file)
    with open(save_question_path, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {save_question_path} successfully! (len: {len(coco_questions_ls)})")

#    /data/linxi/workspace/CHAIR-metric-standalone/I4_sub2000_gt_objects.json 
def coco_prompt_control_noisy(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None, precision=0.5, recall=0.5):
    assert prompt is not None
    if subset is not None:
        save_question_file = f'I{PROMPT_TYPE}_sub{len(subset)}'
    else:
        save_question_file = f'I{PROMPT_TYPE}'
    if control is not None:
        save_question_file += '_control_noisy'
    save_question_file += '.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(path, coco_seg_file)
    coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]
    
    label_dict_path = "./data/category_dict.json"
    category_dict = json.load(open(label_dict_path, "r"))
    all_obj_ls = list(category_dict.values())

    def add_noise(obj_ls, all_obj_ls, precision, recall):
        # add noise from all_obj_ls to the object list
        # precision: the probability of the object in the list is correct
        # recall: the probability of the object in the image is in the list
        # return: noisy object list
        import random
        noisy_obj_ls = []
        for obj in obj_ls:
            if random.random() < recall:
                noisy_obj_ls.append(obj)
        FP_num = int((1/precision-1)*len(noisy_obj_ls)) 
        for i in range(FP_num):      
            noisy_obj_ls.append(random.choice(all_obj_ls))
        return noisy_obj_ls
    
    coco_questions_ls = []
    for i in range(0, len(coco_data), 6):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)
        
        if control is not None:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": control + obj_ls2str(add_noise(coco_seg_data[j]['objects'], all_obj_ls, precision, recall))}]}
        else:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}]}
        coco_questions_ls.append(question)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))

    save_question_path = os.path.join(save_path, "question", save_question_file)
    with open(save_question_path, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {save_question_path} successfully! (len: {len(coco_questions_ls)})")
    
def coco_pope_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None, positive=True):
    assert prompt is not None
    detr_pred_file = "I4_pred_obj_detr_th0.9.json"
    
    if subset is not None:
        save_question_file = f'pope_sub{len(subset)}'
    else:
        save_question_file = 'pope'
    if control is not None:
        save_question_file += '_mmc4_grey' + '_'+ detr_pred_file.split('_')[-1].split('.json')[0]
    save_question_file += '.json'
    save_label_file = save_question_file[:-5] + '_label.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(save_path, "obj", detr_pred_file)
    coco_seg_data = json.load(open(coco_seg_path, 'r'))
    
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]

    coco_questions_ls = []
    coco_labels_ls = []
    for i in range(0, len(coco_data), 1):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)

        prompt = coco_data[i]['text']
        label = coco_data[i]['label']
        # change prompt to negative expression

        # original prompt: Is there a bag in the image?
        # new prompt: Doesn't a bag exist in the image?
        # change the original prompt to new prompt
        if positive is False:
            prompt = prompt[0].upper() + prompt[1:]
            prompt = prompt.replace('Is there', "Doesn't")
            prompt = prompt.replace(' in', ' exist in')
            label = 'No' if label == 'Yes' or label == 'yes' else 'Yes'

        if control is not None:
            # control + obj_ls2str(coco_seg_data[j]['objects']
            neg_prompt = control_pre + obj_ls2str(coco_seg_data[j]['objects']) + " " + control_post + prompt.lower()
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": neg_prompt}]}
            label = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "label": label}
        else:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}]}
            label = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "label": label}
        coco_questions_ls.append(question)
        coco_labels_ls.append(label)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))
    if not os.path.exists(os.path.join(save_path, "label")):
        os.makedirs(os.path.join(save_path, "label"))
    
    question_file = os.path.join(save_path, "question", save_question_file)
    answer_file = os.path.join(save_path, "label", save_label_file)
    with open(question_file, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {question_file} successfully! (len: {len(coco_questions_ls)})")
    with open(answer_file, 'w') as f:
        json.dump(coco_labels_ls, f)
        print(f"save at {answer_file} successfully! (len: {len(coco_labels_ls)})")
    
def coco_pope_control(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset=None, positive=True):
    assert prompt is not None
    if subset is not None:
        save_question_file = f'pope_sub{len(subset)}'
    else:
        save_question_file = 'pope'
    if control is not None:
        save_question_file += '_c3_new'
    if positive is False:
        save_question_file += '_neg'
    save_question_file += '.json'
    save_label_file = save_question_file[:-5] + '_label.json'

    coco_pope_path = os.path.join(path, coco_pope_file)
    coco_data = [json.loads(q) for q in open(coco_pope_path, 'r')]

    coco_seg_path = os.path.join(path, coco_seg_file)
    coco_seg_data = [json.loads(q) for q in open(coco_seg_path, 'r')]
    
    print(len(coco_data), coco_data[0])

    if subset is not None:
        coco_data = [coco_data[i] for i in subset]

    coco_questions_ls = []
    coco_labels_ls = []
    for i in range(0, len(coco_data), 1):
        # find the corresponding image in coco_seg_data, they have the same 'image'
        j = i//6
        while coco_data[i]['image'] != coco_seg_data[j]['image']:
            j = (j+1)%len(coco_seg_data)

        prompt = coco_data[i]['text']
        label = coco_data[i]['label']
        # change prompt to negative expression

        # original prompt: Is there a bag in the image?
        # new prompt: Doesn't a bag exist in the image?
        # change the original prompt to new prompt
        if positive is False:
            prompt = prompt[0].upper() + prompt[1:]
            prompt = prompt.replace('Is there', "Doesn't")
            prompt = prompt.replace(' in', ' exist in')
            label = 'No' if label == 'Yes' or label == 'yes' else 'Yes'

        if control is not None:
            # control + obj_ls2str(coco_seg_data[j]['objects']
            neg_prompt = control_pre + obj_ls2str(coco_seg_data[j]['objects']) + " " + control_post + prompt.lower()
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": neg_prompt}]}
            label = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "label": label}
        else:
            question = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "conversations": [{"from":"human", "value": prompt},{"from":"gpt", "value": "None"}]}
            label = {"id": coco_data[i]['question_id'], "image": coco_data[i]['image'], "label": label}
        coco_questions_ls.append(question)
        coco_labels_ls.append(label)

    # confirm the save path
    if not os.path.exists(os.path.join(save_path, "question")):
        os.makedirs(os.path.join(save_path, "question"))
    if not os.path.exists(os.path.join(save_path, "label")):
        os.makedirs(os.path.join(save_path, "label"))
    
    question_file = os.path.join(save_path, "question", save_question_file)
    answer_file = os.path.join(save_path, "label", save_label_file)
    with open(question_file, 'w') as f:
        json.dump(coco_questions_ls, f)
        print(f"save at {question_file} successfully! (len: {len(coco_questions_ls)})")
    with open(answer_file, 'w') as f:
        json.dump(coco_labels_ls, f)
        print(f"save at {answer_file} successfully! (len: {len(coco_labels_ls)})")

if __name__ == "__main__":
    global PROMPT_TYPE
    PROMPT_TYPE = 4
    path = "../POPE/output/coco"
    # coco_pope_file = "coco_ad_new.json" # zidong
    coco_pope_file = "coco_pope_adversarial.json"
    coco_seg_file = "coco_ground_truth_segmentation.json"

    save_path = "../POPE/llava_qa"
    if PROMPT_TYPE == 3:
        prompt = "Provide a brief description of the given image."
    elif PROMPT_TYPE == 4:
        prompt = "Generate a short caption of the image."
    elif PROMPT_TYPE == 5:
        prompt = "Provide a brief description of the given image."
    else:
        raise NotImplementedError
    # control = "It is confirmed that the following objects appear in the image: " 
    # control = "It has been verified that the image exclusively contains the following objects: "
    control_pre = "This image contains "
    control_post = "Based on this, "
    # control_post = "According to the hint and the image, generate a short caption of the image."
    # control = "Generate a short caption for this image, focusing on the visible elements: "
    # control = ["generate a short caption for this image.", "Focusing on the visible elements in this image: "]
    control = [prompt.lower(), "Focusing on the visible elements in this image: "]
    # control_pope = 
# This image contains a keyboard, laptop, and dining table. Based on this, is there a keyboard visible?

    # subset = [i for i in range(0, 240)]
    subset = None
    # precision = 0.59
    # recall = 0.90
    # control = None
    # coco_prompt_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    coco_qa90_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    # coco_prompt_control(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    # coco_pope_control(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    # coco_pope_control_grey(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    # coco_obj_ls(path, save_path, coco_pope_file, coco_seg_file, prompt, control, subset)
    