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/lucy7-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="pan,box_zoom,wheel_zoom,zoom_in,zoom_out,save,reset,hover", 
           width=1500, height=1000, tooltips="""
    <div>
        <div><strong>@label</strong></div>
        <div><img src="@image" ></div>
    </div>
""")
p.circle('x', 'y', size=5, 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

# Add custom CSS to adjust layout
st.markdown("""
    <style>
        .bokeh-chart-container {
            margin-left: -100px;  /* Adjust the value as needed */
        }
    </style>
    <div class="bokeh-chart-container">
""", unsafe_allow_html=True)

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

# Close the div
st.markdown("</div>", unsafe_allow_html=True)