import os
import openai
import json

import itertools

openai.api_key = '***' #FILL IN YOUR OWN HERE

def stringtolist(description):
    return [descriptor[2:] for descriptor in description.split('\n') if (descriptor != '') and (descriptor.startswith('- '))]

def generate_prompt(category_name: str):
    # you can replace the examples with whatever you want; these were random and worked, could be improved
    return f"""Q: What are useful visual features for distinguishing a lemur in a photo?
A: There are several useful visual features to tell there is a lemur in a photo:
- four-limbed primate
- black, grey, white, brown, or red-brown
- wet and hairless nose with curved nostrils
- long tail
- large eyes
- furry bodies
- clawed hands and feet

Q: What are useful visual features for distinguishing a television in a photo?
A: There are several useful visual features to tell there is a television in a photo:
- electronic device
- black or grey
- a large, rectangular screen
- a stand or mount to support the screen
- one or more speakers
- a power cord
- input ports for connecting to other devices
- a remote control

Q: What are useful features for distinguishing a {category_name} in a photo?
A: There are several useful visual features to tell there is a {category_name} in a photo:
-
"""




def generate_prompt_synonyms(category_name: str):
    return f"""What is the similar words of School bus?
A: There are several words similar with School bus:
- School transport
- Yellow bus
- School coach
- Student bus
- Educational bus
- Pupil transport
- Children's bus
- School vehicle
- Trolly bus


Q: What is the similar words of Television?
A: There are several words similar with Television:
- TV
- Telly
- Tube
- Boob tube
- Small screen
- Idiot box
- Cathode-ray tube
- Vid
- Telly
- Receiver

Q: What are the synonyms of {category_name}?
A: There are several synonyms of {category_name}:
- """



# def generate_prompt_synonyms(category_name: str):
#     return f"""Q: What are the synonyms of School bus?
# A: There are several synonyms of School bus:
# - School transport
# - Yellow bus
# - School coach
# - Student bus
# - Educational bus
# - Pupil transport
# - Children's bus
# - School vehicle


# Q: What are the synonyms of Television?
# A: There are several synonyms of Television:
# - TV
# - Telly
# - Tube
# - Boob tube
# - Small screen
# - Idiot box
# - Cathode-ray tube

# Q: What are the synonyms of {category_name}?
# A: There are several synonyms of {category_name}:
# - """




# generator 
def partition(lst, size):
    for i in range(0, len(lst), size):
        yield list(itertools.islice(lst, i, i + size))
        
def obtain_descriptors_and_save(filename, class_list):
    responses = {}
    descriptors = {}
    
    
    
    prompts = [generate_prompt_synonyms(category.replace('_', ' ')) for category in class_list]
    for entity in prompts:
        print(entity)
        break

    
    # most efficient way is to partition all prompts into the max size that can be concurrently queried from the OpenAI API
    responses = [openai.Completion.create(model="gpt-3.5-turbo-instruct",
    # responses = [openai.Completion.create(model="text-davinci-003",
                                            prompt=prompt_partition,
                                            temperature=0.,
                                            max_tokens=100,
                                            ) for prompt_partition in partition(prompts, 20)]
    response_texts = [r["text"] for resp in responses for r in resp['choices']]
    descriptors_list = [stringtolist(response_text) for response_text in response_texts]
    descriptors = {cat: descr for cat, descr in zip(class_list, descriptors_list)}
    
    # save descriptors to json file
    if not filename.endswith('.json'):
        filename += '.json'
    with open(filename, 'w') as fp:
        json.dump(descriptors, fp, indent=2)
    

# obtain_descriptors_and_save('Example', ['car', 'bus', 'taxi', 'motocycle'])

# file_name='Caltech101'
# file_name='EuroSAT'
# file_name='DescribableTextures'
# file_name='FGVCAircraft'
# file_name='OxfordFlowers'
# file_name='OxfordPets'
# file_name='StanfordCars'
# file_name='SUN397'
# file_name='Food101'
file_name='UCF101'
# file_name='ImageNet'


with open(f'{file_name}.json', 'r') as fp:
    label_list = json.load(fp)

label_list = [*label_list]
obtain_descriptors_and_save(f'{file_name}_syn', label_list)
