# read file in ./POPE/lure
import os
import json
import numpy as np
import yaml
yaml_file="./POPE/lure/lvlm_lure.yaml"
new_dir = "./POPE/lure_new_sent"
# files = os.listdir("./POPE/lure")
os.makedirs(new_dir, exist_ok=True)
# files = [file for file in files if "jsonl" in file]
with open(yaml_file, "r") as f:
    answers_yaml = yaml.load(f, Loader=yaml.FullLoader)
for k, v in answers_yaml["results"]["chair"].items():
    answer_file = v["answer_file"]
    lure_file = v["lure_file"]
    print(f"evaluating {k}...")

    with open(lure_file, "r") as f:
        try:
            data = [json.loads(jline) for jline in f.readlines()]
        except:
            data = json.load(f)
    with open(answer_file, "r") as f:
        if k in ["llava", "mplug-owl2", "llava2"]:
            answers = [json.loads(jline) for jline in f.readlines()]
        else:
            answers = json.load(f)
        key = "text" if "text" in answers[0].keys() else "answer"
        for i in range(len(data)):
            if "caption" in data[i].keys():
                # data[i]["caption"] = data[i]["caption"].split(".")[0]
                
                # cut the caption in data to the same length as answers
                # len0 = len(data[i]["caption"].split())
            
                # data[i]["caption"] = data[i]["caption"].split()[:len(answers[i][key].split())]
                # data[i]["caption"] = " ".join(data[i]["caption"])
                
                # cut to same num of sentences as answers
                num_sent = len(answers[i][key].split("."))
                data[i]["caption"] = ".".join(data[i]["caption"].split(".")[:num_sent])
                
    # write data to new file
    with open(new_dir + "/" + lure_file.split("/")[-1], "w") as f:
        json.dump(data, f)
