import base64
import requests

type = "qa90"
# eval_len = [1]
eval_len = [i for i in range(50,90)]
print(len(eval_len))  

# OpenAI API Key
# new api key
# api_key = "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"
api_key = "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"
# api_key = "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" # hjx

headers = {
  "Content-Type": "application/json",
  "Authorization": f"Bearer {api_key}"
}

# Function to encode the image
def encode_image(image_path):
  with open(image_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')



import json
import os

template_file = "/data/linxi/workspace/gpt4v/template.json"
with open(template_file, "r") as f:
    template = json.load(f)

answers_file = "/data/linxi/workspace/gpt4v/lvlm_answers.yaml"
# read answer file
import yaml
with open(answers_file, "r") as f:
  answers = yaml.load(f, Loader=yaml.FullLoader)

for k, v in answers["results"][type].items():
  print(f"evaluating {k}...")
  answer_file0 = v["answer_file0"]
  answer_file1 = v["answer_file1"]
  
  with open(answer_file1, "r") as f:
    if k in ["llava", "mplug-owl2", "llava2"]:
      data1 = [json.loads(jline) for jline in f.readlines()]
    else:
      data1 = json.load(f)
  with open(answer_file0, "r") as f:
    if k in ["llava", "mplug-owl2", "llava2"]:
      data0 = [json.loads(jline) for jline in f.readlines()]
    else:
      print(answer_file0)
      data0 = json.load(f)
  print(answer_file0,len(data0))
  if "qa90" in type:
    question_file = "/data/linxi/workspace/POPE/llava_qa/question/qa4_mmc4_grey_th0.9.json"
  elif "grey" in type:
    question_file = "/data/linxi/workspace/POPE/llava_qa/question/I4_mmc4.json"
  elif "test" in type:
    question_file = "/data/linxi/workspace/POPE/llava_qa/question/qa4_test.json"
  with open(question_file, "r") as f:
    try:
      question_data = json.load(f)
    except:
      question_data = [json.loads(jline) for jline in f.readlines()]

  # write response to file 
  model = answer_file1.split("/")[4] + "_" + answer_file1.split("/")[-1].split(".json")[0]
  save_path = f"/data/linxi/workspace/gpt4v/results_{type}_0518/{model}.json"
  # i = 1
  question_keys = ["prompt", "question"]
  question_key = question_keys[0] if question_keys[0] in data0[0].keys() else question_keys[1]
  answer_keys = ["text", "answer"]
  answer_key = answer_keys[0] if answer_keys[0] in data0[0].keys() else answer_keys[1]
  for i in eval_len:
    answer0 = data0[i][answer_key]
    answer1 = data1[i][answer_key]
    question = data0[i][question_key]
    image_file = question_data[i]["image"]
    # TODO: change image_dir

    image_dir = "/data/linxi/workspace/POPE/data/val/val2014"
    # image_dir = "/data/linxi/workspace/POPE/data/minival2014/minival2014"
    # image_file = "COCO_val2014_000000144305.jpg"
    image_path = image_dir + "/" + image_file

    # Getting the base64 string
    base64_image = encode_image(image_path)

    payload = {
      "model": "gpt-4-vision-preview",
      "messages": [
        {
            "role": "system",
            "content": [
              {
                "type": "text",
                "text": template["system"]
              }
            ]
        },
        {
          "role": "user",
          "content": [
            {
              "type": "text",
              "text": template["prompt"] + template["question"] + question + "\n"+template["assistant"][0] + answer0 + "\n" + template["assistant"][1] + answer1 +"\n"+ "\n"+ "Output format:\n"+template["output_format"]
            },
            {
              "type": "image_url",
              "image_url": {
                "url": f"data:image/jpeg;base64,{base64_image}"
              }
            }
          ]
        }
      ],
      "max_tokens": 300
    }
    # print(payload["messages"][0]["content"][0]["text"])
    # print(payload['messages'][1]['content'][0]['text'])

    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)

    response = response.json()
    response['question_id'] = question_data[i]["id"]
    response['image'] = question_data[i]["image"]
    print(response)
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    with open(save_path, "a+") as f:
      json.dump(response, f, indent=4)
    print(f"saved to {save_path}")