import torch
from huggingface_hub import hf_hub_download, HfApi
import random


"""
Returns the list of directions currently available in the HF library
"""
def get_all_directions_names():
    hf_api = HfApi()
    info = hf_api.list_models(author="pix2pix-zero-library")
    l_model_ids = [m.modelId for m in info]
    l_model_ids = [m for m in l_model_ids if "_sd14" in m]
    l_edit_names = [m.split("/")[-1] for m in l_model_ids]
    # l_edit_names = [m for m in l_edit_names if "_sd14" in m]

    # pdb.set_trace()
    l_desc = [hf_hub_download(repo_id=m_id, filename="short_description.txt") for m_id in l_model_ids]
    d_name2desc = {k: open(m).read() for k,m in zip(l_edit_names, l_desc)}

    return d_name2desc


def get_emb(dir_name):
    REPO_ID = f"pix2pix-zero-library/{dir_name.replace('.pt','')}"
    if "_sd14" not in REPO_ID: REPO_ID += "_sd14"
    FILENAME = dir_name
    if "_sd14" not in FILENAME: FILENAME += "_sd14"
    if ".pt" not in FILENAME: FILENAME += ".pt"
    ret = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    return torch.load(ret, map_location=torch.device('cpu'), weights_only=True)


def generate_image_prompts_with_templates(word):
    prompts = []
    adjectives = ['majestic', 'cute', 'colorful', 'ferocious', 'elegant', 'graceful', 'slimy', 'adorable', 'scary', 'fuzzy', 'tiny', 'gigantic', 'brave', 'fierce', 'mysterious', 'curious', 'fascinating', 'charming', 'gleaming', 'rare']
    verbs = ['strolling', 'jumping', 'lounging', 'flying', 'sleeping', 'eating', 'playing', 'working', 'gazing', 'standing']
    adverbs = ['gracefully', 'playfully', 'elegantly', 'fiercely', 'curiously', 'fascinatingly', 'charmingly', 'gently', 'slowly', 'quickly', 'awkwardly', 'carelessly', 'cautiously', 'innocently', 'powerfully', 'grumpily', 'mysteriously']
    backgrounds = ['a sunny beach', 'a bustling city', 'a quiet forest', 'a cozy living room', 'a futuristic space station', 'a medieval castle', 'an enchanted forest', 'a misty graveyard', 'a snowy mountain peak', 'a crowded market']

    sentence_structures = {
        "subject verb background": lambda word, bg, verb, adj, adv: f"A {word} {verb} {bg}.",
        "background subject verb": lambda word, bg, verb, adj, adv: f"{bg}, a {word} is {verb}.",
        "adjective subject verb background": lambda word, bg, verb, adj, adv: f"A {adj} {word} is {verb} {bg}.",
        "subject verb adverb background": lambda word, bg, verb, adj, adv: f"A {word} is {verb} {adv} {bg}.",
        "adverb subject verb background": lambda word, bg, verb, adj, adv: f"{adv.capitalize()}, a {word} is {verb} {bg}.",
        "background adjective subject verb": lambda word, bg, verb, adj, adv: f"{bg}, there is a {adj} {word} {verb}.",
        "subject verb adjective background": lambda word, bg, verb, adj, adv: f"A {word} {verb} {adj} {bg}.",
        "adjective subject verb": lambda word, bg, verb, adj, adv: f"A {adj} {word} is {verb}.",
        "subject adjective verb background": lambda word, bg, verb, adj, adv: f"A {word} is {adj} and {verb} {bg}.",
    }

    sentences = []
    for bg in backgrounds:
        for verb in verbs:
            for adj in adjectives:
                adv = random.choice(adverbs)
                sentence = f"A {adv} {adj} {word} {verb} on {bg}."
                sentence_structure = random.choice(list(sentence_structures.keys()))
                sentence = sentence_structures[sentence_structure](word, bg, verb, adj, adv)
                sentences.append(sentence)
    return sentences




if __name__=="__main__":
    print(get_all_directions_names())
    # print(get_emb("dog_sd14.pt").shape)
    # print(get_emb("dog").shape)
    # print(generate_image_prompts("dog")[0:5])