import streamlit as st
import pandas as pd
from bokeh.plotting import figure
from bokeh.models import ColumnDataSource
from bokeh.palettes import Category10

# --- Define Styles ---
st.markdown(
    """
    <style>
    .main-title {
        font-size: 50px;
        color: #4CAF50;
        text-align: center;
    }
    .sub-title {
        font-size: 30px;
        color: #555;
    }
    .custom-text {
        font-size: 18px;
        line-height: 1.5;
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.markdown('<h1 class="main-title">Merit Secret Embeddings 🎒📃🏆</h1>', unsafe_allow_html=True)
st.markdown('<h2 class="sub-title">Donut</h2>', unsafe_allow_html=True)
st.markdown(
    """
    <p class="custom-text">
    Explore how Donut perceives real data. 
    </p>
    """,
    unsafe_allow_html=True
)

# Cargar el CSV
df = pd.read_csv("data/data.csv")
unique_labels = df['label'].unique().tolist()

# Definir tooltips para la imagen y la etiqueta
TOOLTIPS = """
<div>
    <div>
        <img src="@img{safe}" style="width:128px; height:auto; float: left; margin: 0px 15px 15px 0px;" alt="@img" border="2"></img>
    </div>
    <div>
        <span style="font-size: 17px; font-weight: bold;">@label</span>
    </div>
</div>
"""

# Crear contenedor para el gráfico
plot_placeholder = st.empty()

# Función que genera y muestra el gráfico a partir de una lista de labels
def render_plot(selected_labels):
    if not selected_labels:
        st.write("No data to display. Please select at least one subset.")
        return

    filtered_data = df[df['label'].isin(selected_labels)]
    p = figure(width=400, height=400, tooltips=TOOLTIPS, title="")
    
    num_labels = len(selected_labels)
    # Ajuste de la paleta
    if num_labels < 3:
        palette = Category10[3][:num_labels]
    elif num_labels in [3, 4, 5, 6, 7, 8, 9, 10]:
        palette = Category10[num_labels]
    else:
        palette = Category10[10][:num_labels]

    # Graficar cada label por separado para poder asignar su color y legend_label
    for label, color in zip(selected_labels, palette):
        subset = filtered_data[filtered_data['label'] == label]
        source = ColumnDataSource(data=dict(
            x=subset['x'],
            y=subset['y'],
            label=subset['label'],
            img=subset['img']
        ))
        p.scatter('x', 'y', size=12, source=source, color=color, legend_label=label)
    
    p.legend.title = "Subsets"
    p.legend.location = "top_right"
    p.legend.click_policy = "hide"

    plot_placeholder.bokeh_chart(p)

# Mostrar inicialmente el gráfico con todas las etiquetas
render_plot(unique_labels)

# --- Desplegable (multiselect) colocado debajo del gráfico ---
selected_labels = st.multiselect(
    "Select Subsets to Visualize:",
    options=unique_labels,
    default=unique_labels
)

# Al cambiar la selección, volver a renderizar el gráfico
render_plot(selected_labels)