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





# 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(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",
                                            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)
    


file_name='UCF101'

label_list = [
"Apply Eye Makeup",
"Apply Lipstick",
"Archery",
"Baby Crawling",
"Balance Beam",
"Band Marching",
"Baseball Pitch",
"Basketball",
"Basketball Dunk",
"Bench Press",
"Biking",
"Billiards",
"Blow Dry Hair",
"Blowing Candles",
"Body Weight Squats",
"Bowling",
"Boxing Punching Bag",
"Boxing Speed Bag",
"Breast Stroke",
"Brushing Teeth",
"Clean And Jerk",
"Cliff Diving",
"Cricket Bowling",
"Cricket Shot",
"Cutting In Kitchen",
"Diving",
"Drumming",
"Fencing",
"Field Hockey Penalty",
"Floor Gymnastics",
"Frisbee Catch",
"Front Crawl",
"Golf Swing",
"Haircut",
"Hammer Throw",
"Hammering",
"Handstand Pushups",
"Handstand Walking",
"Head Massage",
"High Jump",
"Horse Race",
"Horse Riding",
"Hula Hoop",
"Ice Dancing",
"Javelin Throw",
"Juggling Balls",
"Jump Rope",
"Jumping Jack",
"Kayaking",
"Knitting",
"Long Jump",
"Lunges",
"Military Parade",
"Mixing",
"Mopping Floor",
"Nunchucks",
"Parallel Bars",
"Pizza Tossing",
"Playing Cello",
"Playing Daf",
"Playing Dhol",
"Playing Flute",
"Playing Guitar",
"Playing Piano",
"Playing Sitar",
"Playing Tabla",
"Playing Violin",
"Pole Vault",
"Pommel Horse",
"Pull Ups",
"Punch",
"Push Ups",
"Rafting",
"Rock Climbing Indoor",
"Rope Climbing",
"Rowing",
"Salsa Spin",
"Shaving Beard",
"Shotput",
"Skate Boarding",
"Skiing",
"Skijet",
"Sky Diving",
"Soccer Juggling",
"Soccer Penalty",
"Still Rings",
"Sumo Wrestling",
"Surfing",
"Swing",
"Table Tennis Shot",
"Tai Chi",
"Tennis Swing",
"Throw Discus",
"Trampoline Jumping",
"Typing",
"Uneven Bars",
"Volleyball Spiking",
"Walking With Dog",
"Wall Pushups",
"Writing On Board",
"Yo Yo",

]

obtain_descriptors_and_save(f'{file_name}', label_list)
