import streamlit as st
from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, HoverTool, LinearColorMapper, ColorBar, BasicTicker
from bokeh.transform import linear_cmap
from bokeh.palettes import Viridis256
from datasets import load_dataset

# Load the dataset
dataset = load_dataset("tonyassi/revolve1-embeddings-xy")['train']

# Extract data
data = {
    'x': [item['x'] for item in dataset],
    'y': [item['y'] for item in dataset],
    'label': [f"ID: {item['id']}" for item in dataset],
    'image': [item['image_url'] for item in dataset],
    'id': [item['id'] for item in dataset]  # Include 'id' for color mapping
}

source = ColumnDataSource(data=data)

# Create a color mapper with reversed palette
color_mapper = linear_cmap(field_name='id', palette=Viridis256[::-1], low=0, high=len(data['id']))

# Create the figure
p = figure(title="Image Similarity Data Visualization", tools="pan,box_zoom,wheel_zoom,zoom_in,zoom_out,save,reset,hover", active_scroll="wheel_zoom",
           width=1500, height=1000, tooltips="""
    <div>
        <div><strong>@label</strong></div>
        <div><img src="@image" ></div>
    </div>
""")
p.circle('x', 'y', size=9, source=source, color=color_mapper)  # Apply the color mapper

# Add color bar
color_bar = ColorBar(color_mapper=color_mapper['transform'], width=8, location=(0, 0),
                     ticker=BasicTicker(desired_num_ticks=10))
p.add_layout(color_bar, 'right')  # Position the color bar to the right



st.set_page_config(layout='wide')

st.markdown("""
# Image Similarity Data Visualization

by [Tony Assi](https://www.tonyassi.com/)

Images can be previewed on hover. Images position are based on similarity, images close to each other look similar. The colors represent the best seller index. 0 is best seller. The dataset is [tonyassi/revolve1-embeddings-xy](https://huggingface.co/datasets/tonyassi/revolve1-embeddings-xy).
""")

#st.html('<br><br><br><br>')

# Display the Bokeh figure in Streamlit
st.bokeh_chart(p,use_container_width=True)