# install
# pip install openai

import os
import openai
import csv
import json
from PIL import Image

openai.api_key =""

prompt_prefix_1 = """Generate only an informative and nature paragraph based on the given information(a,b,c):\n"""  
prompt_prefix_2 = """a. Image Resolution:  """
prompt_prefix_3 = """/\n b. Image Caption: """
prompt_prefix_4 = """/\n c. Dense Caption:"""
prompt_suffix = """/\n There are some rules:
Show object, color and position.
Use nouns rather than coordinates to show position information of each object.
No more than 7 sentences.
Only use one paragraph.
Do not appear number.
Do not say information that has nothing to do with the picture.
"""

already_download_path = "coco_test_gpt3_5_response_from_dense_ofa_caption_800.json"
output_data = []
with open(already_download_path , 'r') as f1:
    output_data = json.load(f1)

already_download = len(output_data)

def read_image_width_height(image_path):
    image = Image.open(image_path)
    width, height = image.size
    return width, height

def paragraph_summary_with_gpt(question):
    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo", 
        messages = [
        {"role": "user", "content" : question}]
    )
    return completion['choices'][0]['message']['content']


dense_caption_w_ofa_path = "/annotation/coco_karpathy_test_dense_caption_w_ofa_caption.json"
data_root = "COCO2014/"

with open(dense_caption_w_ofa_path , 'r') as f1:
    data1 = json.load(f1)

num_samples = len(data1)
# num_samples = 802

for i in range(already_download, num_samples):
    width, height = read_image_width_height(data_root + data1[i]['image'])
    instruction = prompt_prefix_1 + prompt_prefix_2 + str(width) + 'X' + str(height)  + prompt_prefix_3 + data1[i]['ofa_caption'] + prompt_prefix_4 + data1[i]['dense_caption'] + prompt_suffix
    gpt_3_5_summary = paragraph_summary_with_gpt(instruction)
    # Create a dictionary for each row
    row_data = {
        'instruction': instruction,
        'input': " ",
        'image': data1[i]['image'],
        'output': gpt_3_5_summary
    }
    # Add the dictionary to the output_data list
    output_data.append(row_data)
    print("{}/{} finished, you have paid OPENAI ${}".format(i, num_samples, i*0.0008125))
    if i % 10 == 0 and i > 0:
        with open('coco_test_gpt3_5_response_from_dense_ofa_caption_{}.json'.format(i), 'w') as outfile:
            json.dump(output_data, outfile, indent=2)

# Save the output_data list to a JSON file

with open('coco_test_gpt3_5_response_from_dense_ofa_caption_{}.json'.format(num_samples), 'w') as outfile:
    json.dump(output_data, outfile, indent=2)