from turtle import color, onclick
import streamlit as st
from PIL import Image, ImageOps
import glob
import json
import requests
import random
import io

if 'index' not in st.session_state:
    st.session_state.index = 0

def set_back():
    st.session_state.index = 0

st.set_page_config(layout="wide")
st.markdown("**This is a demo of the *ImageCoDe* dataset. Sample an example description with the left +/- button on the right and compare all the images with the +/- button on the right. If you want to know the groundtruth solution, scroll down to the end of the page!**")

col1, col2 = st.columns(2)

prefix = 'https://raw.githubusercontent.com/BennoKrojer/imagecode-val-set/main/image-sets-val/'
set2ids = json.load(open('set2ids.json', 'r'))
descriptions = json.load(open('valid_list.json', 'r'))

st.session_state.index = int(col2.number_input('Image Index from 0 to 9', value=0, min_value=0, max_value=9))
example_idx = int(col1.number_input('Sample an example from the validation set', value=0, min_value=0, max_value=len(descriptions)-1, on_change=set_back))
img_set, idx, descr = descriptions[example_idx]
idx = int(idx)
images = [prefix+'/'+img_set+'/'+i for i in set2ids[img_set]]
img_urls = images.copy()

col1.markdown(f'**Description**:')
col1.markdown(descr)

img = images[st.session_state.index]
images[st.session_state.index] = ImageOps.expand(Image.open(io.BytesIO(requests.get(images[st.session_state.index], stream=True).content)),border=20,fill='blue')

caps = list(range(10))
cap = str(st.session_state.index)

col1.image(img, use_column_width=True, caption=cap)
col2.image(images, width=175, caption=caps)

st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown('#')
st.markdown(f'**The correct target image index is: {idx}**')