import os
import openai
import json
from tqdm import tqdm
import time



def split_dict(d, n):
    if n <= 1:
        return [d]
    # Determine the approximate size of each chunk
    chunk_size = len(d) // n
    
    # Split the dictionary into keys
    keys = list(d.keys())
    
    # Split the keys into chunks
    key_chunks = [keys[i:i+chunk_size] for i in range(0, len(keys), chunk_size)]
    
    # Create sub-dictionaries for each chunk
    sub_dicts = [{k: d[k] for k in chunk} for chunk in key_chunks]
    
    return sub_dicts
    
api_key = "<OPENAI_API_KEY>"  # Replace with your own API key
openai.api_key = api_key
load_query = 'imagenet/imagenet_prompt.json'
with open('imagenet/imagenet.json') as f:
    data = json.load(f)

if os.path.exists(load_query):
    with open(load_query) as f:
        Result = json.load(f)
else:
    Result= {}

for (k,full_v) in tqdm(data.items()):
    print(k)
    if k in Result.keys():
        continue
    # Split the dictionary into 5 parts
    chunks = split_dict(full_v, 1)
    c_results = dict()
    for i, v in enumerate(chunks):
        print("Chunk:", i)
        q = "{}:{}".format(k,json.dumps(v, indent=2))
        # Set the prompt text to generate a response for
        prompt_text = '''Q: Generate the all sentences that describe a concept according to the each attribute, as json 
"beaver": {
    "Concepts": [
        "Animal",
        "Mammal",
        "Semi-aquatic",
        "Herbivore"
    ],
    "Substance": [
        "four legs",
        "two ears",
        "two eyes",
        "a nose",
        "a flat tail",
        "fur"
    ],
    "Attributes": {
        "Head": {
            "shape": "rounded",
            "Nose": [
                "moist",
                "black"
            ],
            "Ears": [
                "small",
                "not easily visible"
            ],
            "Eyes": [
                "small",
                "dark colored"
            ]
        },
        "Body": {
            "description": [
                "covered in fur",
                "four-legged",
                "cylindrical and muscular"
            ],
            "Tail": [
                "flat",
                "wide",
                "used for swimming and slapping water"
            ],
            "Legs": [
                "short",
                "powerful",
                "webbed for swimming"
            ]
        }
    },
    "Environment": [
        "rivers",
        "lakes",
        "streams"
    ]
}
A: 
{
    "Concepts": [
        "A beaver is an animal.",
        "A beaver is a mammal.",
        "A beaver is semi-aquatic.",
        "A beaver is a herbivore."
    ],
    "Substance": [
        "A beaver has four legs.",
        "A beaver has two ears.",
        "A beaver has two eyes.",
        "A beaver has a nose.",
        "A beaver has a flat tail.",
        "A beaver has fur."
    ],
    "Attributes":{
        "Head": {
            "shape": "A beaver's head is rounded in shape.",
            "Nose": [
                "A beaver's nose is moist.",
                "A beaver's nose is black."
            ],
            "Ears": [
                "A beaver's ears are small.",
                "A beaver's ears are not easily visible."
            ],
            "Eyes": [
                "A beaver's eyes are small.",
                "A beaver's eyes are dark colored."
            ]
        },
        "Body": {
            "description": [
                "A beaver's body is covered in fur.",
                "A beaver's body is four-legged.",
                "A beaver's body is cylindrical and muscular."
            ],
            "Tail": [
                "A beaver's tail is flat.",
                "A beaver's tail is wide.",
                "A beaver's tail is used for swimming and slapping water."
            ],
            "Legs": [
                "A beaver's legs are short.",
                "A beaver's legs are powerful.",
                "A beaver's legs are webbed for swimming."
            ]
        }
    },
    "Environment": [
        "A beaver lives in rivers.",
        "A beaver lives in lakes.",
        "A beaver lives in streams."
    ]
}
Q: Generate the all sentences that describe a concept according to the each attribute, as json  
'''
        prompt_text += q
        prompt_text += '\nA:'
        result = ""

    

        prompt = [{"role": "user", "content": prompt_text}]

        cur_res = None
        while cur_res is None:
            try:
                completion = openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=prompt
                )

                res = completion.choices[0].message["content"]
                cur_res = json.loads(res)
            except Exception as e: 
                print(f"Error: {e}")
                print(f"Sleeping for 10 seconds...")
                time.sleep(10)


        
        Result[k] = cur_res

    

    with open(load_query, "w") as f:
        json.dump(Result, f, indent=4)