import pandas
import streamlit as st
from diffusers import StableDiffusionPipeline

from game_manager import GameManager

game_manager = GameManager()


def empty_cache():
    st.empty()


# @st.cache_resource  # 👈 Add the caching decorator
def get_pipe():
    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cpu")
    # pipe.safety_checker = lambda images, clip_input: (images, False)
    return pipe


def start_game():
    INFERENCE_STEPS = 15
    PROMPT_STRENGTH = 10.0
    HEIGHT = 128  # must be divisible by 8
    WIDTH = 128  # must be divisible by 8
    prompt_base = ""
    if nude_selected:
        prompt_base = "Realisitc photograph of naked Elon Musk, style "
    else:
        prompt_base = "Realisitc photograph of Elon Musk, style "
    prompt = prompt_base + game_manager.get_artist() + " professions, RAW photo, *subject*, high detailed skin:2.3, 8k uhd, dslr, studio lighting, studio lighting, Fujifilm XT3"
    pipe = get_pipe()
    image = pipe(prompt, height=HEIGHT, width=WIDTH, num_inference_steps=INFERENCE_STEPS,
                 guidance_scale=PROMPT_STRENGTH).images[0]
    st.image(image)


styles = pandas.read_csv("./artist_styles")
st.title('AI Draws :blue[Elon Musk] :sunglasses:')
select = st.selectbox('Choose A Style', styles["Artist"], key='1', index=140)
st.button("Draw Elon Musk (maybe NSFW)", on_click=start_game)
nude_selected = st.checkbox("Nude")
artist_select = st.text("Style Selected:" + select)
if select:
    artist_select.text = "Style Selected: " + select
    game_manager.set_artist(select)
st.text("Wait Time:  This is a slow CPU, wait for 2 minutes for images")
st.text("Note: you can select styles by typing")
st.text("Note: to save the image right-click then select 'save image as'")
st.title('With more iterations this would actually look like :blue[Elon Musk] :sunglasses:')

st.button("Empty Cache", on_click=empty_cache)