import requests
# Function to query an API with a given text

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

def merge_dicts(dict1, dict2):
    """
    Recursively merge two dictionaries.
    """
    result = dict1.copy()
    for key, value in dict2.items():
        if key in result:
            if isinstance(result[key], dict) and isinstance(value, dict):
                result[key] = merge_dicts(result[key], value)
            elif isinstance(result[key], list) and isinstance(value, list):
                result[key].extend(list(set(value) - set(result[key])))
            else:
                result[key] = value
        else:
            result[key] = value
    return result

# Function to read even lines from a file
def read_lines(filename):
    with open(filename, 'r') as f:
        lines = f.read().splitlines()
    return lines

load_query = 'imagenet/imagenet.json'
api_key = "<OPENAI_API_KEY>"  # Replace with your own API key
openai.api_key = api_key
# Example usage
if os.path.exists(load_query):
    with open(load_query) as f:
        Result = json.load(f)
else:
    Result = dict()
filename = 'imagenet/imagenet-classes.txt'
lines = read_lines(filename)
print(len(lines))
print(len(Result.keys()))
print(len(set(lines)))
for cs in tqdm(lines):
    # print(cs)
    if cs in Result.keys():
        continue
    print(cs)
    for c in [cs]:
        # Set the prompt text to generate a response for
        prompt_text = '''Q: this is a dear because
    A:
    {
        "Concepts": [
            "Animal",
            "Mammal",
            "Wild",
            "Herbivore"
        ],
        "Substance": [
            "four legs",
            "two ears",
            "two eyes",
            "a nose",
            "a tail",
            "fur",
            "antlers (in males)"
        ],
        "Attributes": {
            "Head": {
                "shape": "elongated",
                "Nose": [
                    "moist",
                    "black"
                ],
                "Ears": [
                    "pointy",
                    "sensitive"
                ],
                "Eyes": [
                    "large",
                    "brown"
                ]
            },
            "Body": {
                "description": [
                    "covered in fur",
                    "four-legged",
                    "slender"
                ],
                "Tail": [
                    "short",
                    "furry"
                ],
                "Legs": [
                    "long",
                    "slender",
                    "hooves"
                ],
                "Antlers (in males)": [
                    "branched",
                    "bony structures on the head",
                    "used for defense and display"
                ]
            }
        },
        "Environment": [
            "forests",
            "meadows",
            "grasslands"
        ]
    }

    Q: this is a goldfish because    
    A:
{
    "Concepts": [
        "Animal",
        "Fish",
        "Aquatic",
        "Pet",
    ],
    "Substance": [
        "scales",
        "fins",
        "gills",
        "tail",
        "eyes"
    ],
    "Attributes": {
        "Head": {
            "shape": "rounded",
            "Mouth": [
                "protruding",
                "small",
                "toothed",
                "toothless"
            ],
            "Eyes": [
                "large",
                "round",
                "located on sides of head"
            ]
        },
        "Body": {
            "Scale": [
                "small",
                "overlapping"
            ],
            "Fins": [
                "two fins on either side",
                "dorsal fin on back",
                "caudal fin at tail end"
            ],
            "Shape": [
                "rounded body shape"
            ],
            "Color": [
                "gold",
                "orange"
            ]
        }
    },
    "Environment": [
        "aquariums",
        "freshwater or saltwater bodies",
        "coral reefs",
        "lakes",
        "rivers",
        "ponds"
    ]
}
    '''
        prompt_text += '''Q: this is a {} because
    A:
    '''.format(c)
        

        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)

        print(res)
        if cs in Result.keys():
            Result[cs] = merge_dicts(Result[cs], cur_res)
        else:
            Result[cs] = cur_res 

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


