import gradio as gr
from transformers import pipeline

from text_utils import wrap_text, compute_text_position
from gan_utils import load_img_generator, generate_img
from PIL import ImageFont, ImageDraw
import torch

# device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = "cpu"

text_generator = pipeline('text-generation', model='huggingtweets/bestmusiclyric')


def generate_captioned_img(lyrics_prompt, gan_model):
    gan_image = generate_img(device, gan_model)

    generated_text = text_generator(lyrics_prompt)[0]["generated_text"]
    wrapped_text = wrap_text(generated_text)

    text_pos = compute_text_position(wrapped_text)

    # Source: https://stackoverflow.com/a/16377244
    draw = ImageDraw.Draw(gan_image)
    font = ImageFont.truetype("DejaVuSans.ttf", 64)
    draw.text((10, text_pos), text=wrapped_text, fill_color=(255, 255, 255), font=font, stroke_fill=(0, 0, 0),
              stroke_width=5)

    return gan_image


iface = gr.Interface(fn=generate_captioned_img, inputs=[gr.Textbox(value="Running with the wolves", label="Lyrics prompt", lines=1),
                                                        gr.Radio(value="aurora",
                                                            choices=["painting", "fauvism-still-life", "aurora",
                                                                     "universe", "moongate"],
                                                                 label="FastGAN model")
                                                        ],
                     outputs="image",
                     allow_flagging="never",
                     title="Illustrated lyrics generator",
                     description="Combines song lyrics generation via the [Best Music Lyric Bot]"
                                 "(https://huggingface.co/huggingtweets/bestmusiclyric) with an artwork randomly "
                                 "generated by a [FastGAN model](https://huggingface.co/spaces/huggan/FastGan).\n\n"
                                 "Text and lyrics are generated independently. "
                                 "If you can implement this idea with images conditioned on the lyrics,"
                                 " I'd be very interested in seeing that!🤗\n\n"
                                 "At the bottom of the page, you can click some example inputs to get you started.",
                     examples=[["Hey now", "fauvism-still-life"], ["It's gonna take a lot", "universe"],
                               ["Running with the wolves", "aurora"], ["His palms are sweaty", "painting"],
                               ["I just met you", "moongate"]]
                     )
iface.launch()


#examples=[["Hey now", "painting"], ["It's gonna take a lot", "universe"], ["So close", "aurora"], ["I just met you", "moongate"],
#                               ["His palms are sweaty", "aurora"]])