import streamlit as st
from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, HoverTool

from datasets import load_dataset
from bokeh.models import ColumnDataSource, LinearColorMapper, ColorBar
from bokeh.transform import linear_cmap
from bokeh.palettes import Viridis256  # You can choose any palette you like and reverse it using [::-1]
from bokeh.models import BasicTicker

# Load the dataset
dataset = load_dataset("tonyassi/lucy6-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="Scatter Plot with Image Hover", tools="hover", tooltips="""
    <div>
        <div><strong>@label</strong></div>
        <div><img src="@image" ></div>
    </div>
""")
p.circle('x', 'y', size=10, 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.html("""
<br><br><br><br><br><br>
""")

# Display the Bokeh figure in Streamlit
st.bokeh_chart(p)