import csv
import json
from PIL import Image

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.
"""


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

dense_caption_w_ofa_path = "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)

output_data = []

# len(data1)
num_samples = 100

for i in range(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
        
    # Create a dictionary for each row
    row_data = {
        'instruction': instruction,
        'input': " ",
        'image': data1[i]['image'],
        'output': " "
    }

    # Add the dictionary to the output_data list
    output_data.append(row_data)

# Save the output_data list to a JSON file
with open('coco_test_gpt4_response_from_dense_ofa_caption_{}.json'.format(num_samples), 'w') as outfile:
    json.dump(output_data, outfile, indent=2)