import streamlit as st
import io
import sys
import time
import json
sys.path.append("./virtex/")

def gen_show_caption(sub_prompt=None):
    with st.spinner("Generating Caption"):
        subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt)
        st.header("Predicted Caption:\n\n")
        st.subheader(f"r/{subreddit}:\t{caption}\n")
    

st.title("Image Captioning Demo from Redcaps")
st.sidebar.markdown(
    """
    Image Captioning Model from VirTex trained on Redcaps
    """
)

with st.spinner("Loading Model"):
    from model import *
    sample_images = get_samples()
    virtexModel = VirTexModel()
    imageLoader = ImageLoader()
    valid_subs = json.load(open(VALID_SUBREDDITS_PATH))
    valid_subs.insert(0, None)

random_image = get_rand_img(sample_images)

st.sidebar.title("Select a sample image")
sample_image = st.sidebar.selectbox(
    "",
    sample_images
)

if st.sidebar.button("Random Sample Image"):
    random_image = get_rand_img(sample_images)
    sample_image = None
    
uploaded_image = None
with st.sidebar.form("file-uploader-form", clear_on_submit=True):
    uploaded_file = st.file_uploader("Choose a file")
    submitted = st.form_submit_button("Submit")
    if uploaded_file is not None and submitted:
        uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
        
st.sidebar.title("Select a Subreddit")
sub = st.sidebar.selectbox(
    "Select None for a Predicted Subreddit",
    valid_subs
)


if uploaded_image is None and submitted:
    st.write("Please select a file to upload")

else:
    image_file = sample_image if sample_image is not None else random_image
    
    image = uploaded_image if uploaded_image is not None else Image.open(image_file)
    
    image_dict = imageLoader.transform(image)
    
#     image = imageLoader.to_image(image_dict["image"].squeeze(0))
    
    show = st.image(image)
    show.image(image, "Your Image")
    
    gen_show_caption(sub)
        
    image.close()

# from model import *
# download_files()
# sample_images = get_samples()
# v, il = VirTexModel(), ImageLoader()

# for s in sample_images:
#     subreddit, caption = v.predict(il.load(s))
#     print("=====================")
#     print(subreddit)
#     print(caption)