import streamlit as st
import pandas as pd
import numpy as np
from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, DataTable, TableColumn, CustomJS, Select, Button, HoverTool, LinearColorMapper, ColorBar, FuncTickFormatter, FixedTicker
from bokeh.layouts import column
from bokeh.palettes import Reds9, Blues9, Oranges9, Purples9, Greys9, BuGn9, Greens9
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE, trustworthiness
from sklearn.metrics import pairwise_distances
import io
import ot
from sklearn.linear_model import LinearRegression
from scipy.stats import binned_statistic_2d
import json


N_COMPONENTS = 2
TSNE_NEIGHBOURS = 150
# WEIGHT_FACTOR = 0.05

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>
"""

def config_style():
    # st.set_page_config(layout="wide")
    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; }
        .bk-legend {
            max-height: 200px;
            overflow-y: auto;
        }
        </style>
    """, unsafe_allow_html=True)
    st.markdown('<h1 class="main-title">Merit Embeddings 🎒📃🏆</h1>', unsafe_allow_html=True)

def load_embeddings(model, version, embedding_prefix, weight_factor):
    if model == "Donut":
        df_real = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_secret_all_{weight_factor}embeddings.csv")
        df_par = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-paragraph-degradation-seq_{weight_factor}embeddings.csv")
        df_line = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-line-degradation-seq_{weight_factor}embeddings.csv")
        df_seq  = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-seq_{weight_factor}embeddings.csv")
        df_rot  = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-rotation-degradation-seq_{weight_factor}embeddings.csv")
        df_zoom = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-zoom-degradation-seq_{weight_factor}embeddings.csv")
        df_render = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_es-render-seq_{weight_factor}embeddings.csv")
        df_pretratrained = pd.read_csv(f"data/donut/{version}/{embedding_prefix}/de_Rodrigo_merit_aux_IIT-CDIP_{weight_factor}embeddings.csv")
        
        # Asignar etiquetas de versión
        df_real["version"] = "real"
        df_par["version"] = "synthetic"
        df_line["version"] = "synthetic"
        df_seq["version"] = "synthetic"
        df_rot["version"] = "synthetic"
        df_zoom["version"] = "synthetic"
        df_render["version"] = "synthetic"
        df_pretratrained["version"] = "pretrained"

        # Asignar fuente (source)
        df_par["source"] = "es-digital-paragraph-degradation-seq"
        df_line["source"] = "es-digital-line-degradation-seq"
        df_seq["source"] = "es-digital-seq"
        df_rot["source"] = "es-digital-rotation-degradation-seq"
        df_zoom["source"] = "es-digital-zoom-degradation-seq"
        df_render["source"] = "es-render-seq"
        df_pretratrained["source"] = "pretrained"

        return {"real": df_real, 
                "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True),
                "pretrained": df_pretratrained}
    
    elif model == "Idefics2":
        df_real = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_secret_britanico_{weight_factor}embeddings.csv")
        df_par = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-paragraph-degradation-seq_{weight_factor}embeddings.csv")
        df_line = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-line-degradation-seq_{weight_factor}embeddings.csv")
        df_seq  = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-seq_{weight_factor}embeddings.csv")
        df_rot  = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-rotation-degradation-seq_{weight_factor}embeddings.csv")
        df_zoom = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-digital-zoom-degradation-seq_{weight_factor}embeddings.csv")
        df_render = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_es-render-seq_{weight_factor}embeddings.csv")
        
        # Cargar ambos subconjuntos pretrained y combinarlos
        df_pretratrained_PDFA = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_aux_PDFA_{weight_factor}embeddings.csv")
        df_pretratrained_IDL = pd.read_csv(f"data/idefics2/{version}/{embedding_prefix}/de_Rodrigo_merit_aux_IDL_{weight_factor}embeddings.csv")
        df_pretratrained = pd.concat([df_pretratrained_PDFA, df_pretratrained_IDL], ignore_index=True)
        
        # Asignar etiquetas de versión
        df_real["version"] = "real"
        df_par["version"] = "synthetic"
        df_line["version"] = "synthetic"
        df_seq["version"] = "synthetic"
        df_rot["version"] = "synthetic"
        df_zoom["version"] = "synthetic"
        df_render["version"] = "synthetic"
        df_pretratrained["version"] = "pretrained"

        # Asignar fuente (source)
        df_par["source"] = "es-digital-paragraph-degradation-seq"
        df_line["source"] = "es-digital-line-degradation-seq"
        df_seq["source"] = "es-digital-seq"
        df_rot["source"] = "es-digital-rotation-degradation-seq"
        df_zoom["source"] = "es-digital-zoom-degradation-seq"
        df_render["source"] = "es-render-seq"
        df_pretratrained["source"] = "pretrained"

        return {"real": df_real, 
                "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True),
                "pretrained": df_pretratrained}
    
    else:
        st.error("Modelo no reconocido")
        return None

def split_versions(df_combined, reduced):
    # Asignar las coordenadas si la reducción es 2D
    if reduced.shape[1] == 2:
        df_combined['x'] = reduced[:, 0]
        df_combined['y'] = reduced[:, 1]
    df_real = df_combined[df_combined["version"] == "real"].copy()
    df_synth = df_combined[df_combined["version"] == "synthetic"].copy()
    df_pretrained = df_combined[df_combined["version"] == "pretrained"].copy()
    
    unique_real = sorted(df_real['label'].unique().tolist())
    unique_synth = {}
    for source in df_synth["source"].unique():
        unique_synth[source] = sorted(df_synth[df_synth["source"] == source]['label'].unique().tolist())
    unique_pretrained = sorted(df_pretrained['label'].unique().tolist())
    
    df_dict = {"real": df_real, "synthetic": df_synth, "pretrained": df_pretrained}
    unique_subsets = {"real": unique_real, "synthetic": unique_synth, "pretrained": unique_pretrained}
    return df_dict, unique_subsets

def get_embedding_from_df(df):
    # Retorna el embedding completo (4 dimensiones en este caso) guardado en la columna 'embedding'
    if 'embedding' in df.columns:
        return np.stack(df['embedding'].to_numpy())
    elif 'x' in df.columns and 'y' in df.columns:
        return df[['x', 'y']].values
    else:
        raise ValueError("No se encontró embedding o coordenadas x,y en el DataFrame.")

def compute_cluster_distance(synthetic_points, real_points, metric="wasserstein", bins=20):
    if metric.lower() == "wasserstein":
        n = synthetic_points.shape[0]
        m = real_points.shape[0]
        weights = np.ones(n) / n
        weights_real = np.ones(m) / m
        M = ot.dist(synthetic_points, real_points, metric='euclidean')
        return ot.emd2(weights, weights_real, M)
    elif metric.lower() == "euclidean":
        center_syn = np.mean(synthetic_points, axis=0)
        center_real = np.mean(real_points, axis=0)
        return np.linalg.norm(center_syn - center_real)
    elif metric.lower() == "kl":
        # Para KL usamos histogramas multidimensionales con límites globales en cada dimensión
        all_points = np.vstack([synthetic_points, real_points])
        edges = [
            np.linspace(np.min(all_points[:, i]), np.max(all_points[:, i]), bins+1)
            for i in range(all_points.shape[1])
        ]
        H_syn, _ = np.histogramdd(synthetic_points, bins=edges)
        H_real, _ = np.histogramdd(real_points, bins=edges)
        eps = 1e-10
        P = H_syn + eps
        Q = H_real + eps
        P = P / P.sum()
        Q = Q / Q.sum()
        kl = np.sum(P * np.log(P / Q))
        return kl
    else:
        raise ValueError("Métrica desconocida. Usa 'wasserstein', 'euclidean' o 'kl'.")

def compute_cluster_distances_synthetic_individual(synthetic_df: pd.DataFrame, df_real: pd.DataFrame, real_labels: list, metric="wasserstein", bins=20) -> pd.DataFrame:
    distances = {}
    groups = synthetic_df.groupby(['source', 'label'])
    for (source, label), group in groups:
        key = f"{label} ({source})"
        data = get_embedding_from_df(group)
        distances[key] = {}
        for real_label in real_labels:
            real_data = get_embedding_from_df(df_real[df_real['label'] == real_label])
            d = compute_cluster_distance(data, real_data, metric=metric, bins=bins)
            distances[key][real_label] = d
    for source, group in synthetic_df.groupby('source'):
        key = f"Global ({source})"
        data = get_embedding_from_df(group)
        distances[key] = {}
        for real_label in real_labels:
            real_data = get_embedding_from_df(df_real[df_real['label'] == real_label])
            d = compute_cluster_distance(data, real_data, metric=metric, bins=bins)
            distances[key][real_label] = d
    return pd.DataFrame(distances).T

def compute_continuity(X, X_embedded, n_neighbors=5):
    n = X.shape[0]
    D_high = pairwise_distances(X, metric='euclidean')
    D_low = pairwise_distances(X_embedded, metric='euclidean')
    indices_high = np.argsort(D_high, axis=1)
    indices_low = np.argsort(D_low, axis=1)
    k_high = indices_high[:, 1:n_neighbors+1]
    k_low = indices_low[:, 1:n_neighbors+1]
    total = 0.0
    for i in range(n):
        set_high = set(k_high[i])
        set_low = set(k_low[i])
        missing = set_high - set_low
        for j in missing:
            rank = np.where(indices_low[i] == j)[0][0]
            total += (rank - n_neighbors)
    norm = 2.0 / (n * n_neighbors * (2*n - 3*n_neighbors - 1))
    continuity_value = 1 - norm * total
    return continuity_value

def create_table(df_distances):
    df_table = df_distances.copy()
    df_table.reset_index(inplace=True)
    df_table.rename(columns={'index': 'Synthetic'}, inplace=True)
    min_row = {"Synthetic": "Min."}
    mean_row = {"Synthetic": "Mean"}
    max_row = {"Synthetic": "Max."}
    for col in df_table.columns:
        if col != "Synthetic":
            min_row[col] = df_table[col].min()
            mean_row[col] = df_table[col].mean()
            max_row[col] = df_table[col].max()
    df_table = pd.concat([df_table, pd.DataFrame([min_row, mean_row, max_row])], ignore_index=True)
    source_table = ColumnDataSource(df_table)
    columns = [TableColumn(field='Synthetic', title='Synthetic')]
    for col in df_table.columns:
        if col != 'Synthetic':
            columns.append(TableColumn(field=col, title=col))
    total_height = 30 + len(df_table)*28
    data_table = DataTable(source=source_table, columns=columns, sizing_mode='stretch_width', height=total_height)
    return data_table, df_table, source_table

def create_figure(dfs, unique_subsets, color_maps, model_name):
    # Se crea el plot para el embedding reducido (asumiendo que es 2D)
    fig = figure(width=600, height=600, tools="wheel_zoom,pan,reset,save", active_scroll="wheel_zoom", tooltips=TOOLTIPS, title="")
    fig.match_aspect = True
    
    # Renderizar datos reales
    real_renderers = add_dataset_to_fig(fig, dfs["real"], unique_subsets["real"],
                                        marker="circle", color_mapping=color_maps["real"],
                                        group_label="Real")
    
    # Renderizar datos sintéticos (por fuente)
    marker_mapping = {
        "es-digital-paragraph-degradation-seq": "x",
        "es-digital-line-degradation-seq": "cross",
        "es-digital-seq": "triangle",
        "es-digital-rotation-degradation-seq": "diamond",
        "es-digital-zoom-degradation-seq": "asterisk",
        "es-render-seq": "inverted_triangle"
    }
    synthetic_renderers = {}
    synth_df = dfs["synthetic"]
    for source in unique_subsets["synthetic"]:
        df_source = synth_df[synth_df["source"] == source]
        marker = marker_mapping.get(source, "square")
        renderers = add_synthetic_dataset_to_fig(fig, df_source, unique_subsets["synthetic"][source],
                                                  marker=marker,
                                                  color_mapping=color_maps["synthetic"][source],
                                                  group_label=source)
        synthetic_renderers.update(renderers)
    
    # Agregar el subset pretrained (se puede usar un marcador distinto, por ejemplo, "triangle")
    pretrained_renderers = add_dataset_to_fig(fig, dfs["pretrained"], unique_subsets["pretrained"],
                                               marker="triangle", color_mapping=color_maps["pretrained"],
                                               group_label="Pretrained")
    
    fig.legend.location = "top_right"
    fig.legend.click_policy = "hide"
    show_legend = st.checkbox("Show Legend", value=False, key=f"legend_{model_name}")
    fig.legend.visible = show_legend
    return fig, real_renderers, synthetic_renderers, pretrained_renderers

def add_dataset_to_fig(fig, df, selected_labels, marker, color_mapping, group_label):
    renderers = {}
    for label in selected_labels:
        subset = df[df['label'] == label]
        if subset.empty:
            continue
        source = ColumnDataSource(data=dict(
            x=subset['x'],
            y=subset['y'],
            label=subset['label'],
            img=subset.get('img', "")
        ))
        color = color_mapping[label]
        legend_label = f"{label} ({group_label})"
        if marker == "circle":
            r = fig.circle('x', 'y', size=10, source=source,
                           fill_color=color, line_color=color,
                           legend_label=legend_label)
        elif marker == "square":
            r = fig.square('x', 'y', size=10, source=source,
                           fill_color=color, line_color=color,
                           legend_label=legend_label)
        elif marker == "triangle":
            r = fig.triangle('x', 'y', size=12, source=source,
                             fill_color=color, line_color=color,
                             legend_label=legend_label)
        renderers[label + f" ({group_label})"] = r
    return renderers

def add_synthetic_dataset_to_fig(fig, df, labels, marker, color_mapping, group_label):
    renderers = {}
    for label in labels:
        subset = df[df['label'] == label]
        if subset.empty:
            continue
        source_obj = ColumnDataSource(data=dict(
            x=subset['x'],
            y=subset['y'],
            label=subset['label'],
            img=subset.get('img', "")
        ))
        color = color_mapping[label]
        legend_label = group_label
        if marker == "square":
            r = fig.square('x', 'y', size=10, source=source_obj,
                           fill_color=color, line_color=color,
                           legend_label=legend_label)
        elif marker == "triangle":
            r = fig.triangle('x', 'y', size=12, source=source_obj,
                             fill_color=color, line_color=color,
                             legend_label=legend_label)
        elif marker == "inverted_triangle":
            r = fig.inverted_triangle('x', 'y', size=12, source=source_obj,
                                      fill_color=color, line_color=color,
                                      legend_label=legend_label)
        elif marker == "diamond":
            r = fig.diamond('x', 'y', size=10, source=source_obj,
                            fill_color=color, line_color=color,
                            legend_label=legend_label)
        elif marker == "cross":
            r = fig.cross('x', 'y', size=12, source=source_obj,
                          fill_color=color, line_color=color,
                          legend_label=legend_label)
        elif marker == "x":
            r = fig.x('x', 'y', size=12, source=source_obj,
                      fill_color=color, line_color=color,
                      legend_label=legend_label)
        elif marker == "asterisk":
            r = fig.asterisk('x', 'y', size=12, source=source_obj,
                             fill_color=color, line_color=color,
                             legend_label=legend_label)
        else:
            r = fig.circle('x', 'y', size=10, source=source_obj,
                           fill_color=color, line_color=color,
                           legend_label=legend_label)
        renderers[label + f" ({group_label})"] = r
    return renderers

def get_color_maps(unique_subsets):
    color_map = {}
    num_real = len(unique_subsets["real"])
    red_palette = Reds9[:num_real] if num_real <= 9 else (Reds9 * ((num_real // 9) + 1))[:num_real]
    color_map["real"] = {label: red_palette[i] for i, label in enumerate(sorted(unique_subsets["real"]))}
    
    color_map["synthetic"] = {}
    for source, labels in unique_subsets["synthetic"].items():
        if source == "es-digital-seq":
            palette = Blues9[:len(labels)] if len(labels) <= 9 else (Blues9 * ((len(labels)//9)+1))[:len(labels)]
        elif source == "es-digital-line-degradation-seq":
            palette = Purples9[:len(labels)] if len(labels) <= 9 else (Purples9 * ((len(labels)//9)+1))[:len(labels)]
        elif source == "es-digital-paragraph-degradation-seq":
            palette = BuGn9[:len(labels)] if len(labels) <= 9 else (BuGn9 * ((len(labels)//9)+1))[:len(labels)]
        elif source == "es-digital-rotation-degradation-seq":
            palette = Greys9[:len(labels)] if len(labels) <= 9 else (Greys9 * ((len(labels)//9)+1))[:len(labels)]
        elif source == "es-digital-zoom-degradation-seq":
            palette = Oranges9[:len(labels)] if len(labels) <= 9 else (Oranges9 * ((len(labels)//9)+1))[:len(labels)]
        elif source == "es-render-seq":
            palette = Greens9[:len(labels)] if len(labels) <= 9 else (Greens9 * ((len(labels)//9)+1))[:len(labels)]
        else:
            palette = Blues9[:len(labels)] if len(labels) <= 9 else (Blues9 * ((len(labels)//9)+1))[:len(labels)]
        color_map["synthetic"][source] = {label: palette[i] for i, label in enumerate(sorted(labels))}
    
    # Asignar colores al subset pretrained usando, por ejemplo, la paleta Purples9
    num_pretrained = len(unique_subsets["pretrained"])
    purple_palette = Purples9[:num_pretrained] if num_pretrained <= 9 else (Purples9 * ((num_pretrained // 9) + 1))[:num_pretrained]
    color_map["pretrained"] = {label: purple_palette[i] for i, label in enumerate(sorted(unique_subsets["pretrained"]))}
    
    return color_map

def calculate_cluster_centers(df, labels):
    centers = {}
    for label in labels:
        subset = df[df['label'] == label]
        if not subset.empty and 'x' in subset.columns and 'y' in subset.columns:
            centers[label] = (subset['x'].mean(), subset['y'].mean())
    return centers

def compute_global_regression(df_combined, embedding_cols, tsne_params, df_f1, reduction_method="t-SNE", distance_metric="wasserstein"):
    if reduction_method == "PCA":
        reducer = PCA(n_components=N_COMPONENTS)
    else:
        reducer = TSNE(n_components=2, random_state=42, 
                         perplexity=tsne_params["perplexity"], 
                         learning_rate=tsne_params["learning_rate"])
    
    reduced = reducer.fit_transform(df_combined[embedding_cols].values)
    # Guardamos el embedding completo (por ejemplo, 4 dimensiones en PCA)
    df_combined['embedding'] = list(reduced)
    # Si el embedding es 2D, asignamos x e y para visualización
    if reduced.shape[1] == 2:
        df_combined['x'] = reduced[:, 0]
        df_combined['y'] = reduced[:, 1]
    
    explained_variance = None
    if reduction_method == "PCA":
        explained_variance = reducer.explained_variance_ratio_
    
    trust = None
    cont = None
    if reduction_method == "t-SNE":
        X = df_combined[embedding_cols].values
        trust = trustworthiness(X, reduced, n_neighbors=TSNE_NEIGHBOURS)
        cont = compute_continuity(X, reduced, n_neighbors=TSNE_NEIGHBOURS)
    
    dfs_reduced, unique_subsets = split_versions(df_combined, reduced)
    
    df_distances = compute_cluster_distances_synthetic_individual(
        dfs_reduced["synthetic"],
        dfs_reduced["real"],
        unique_subsets["real"],
        metric=distance_metric
    )
    
    global_distances = {}
    for idx in df_distances.index:
        if idx.startswith("Global"):
            source = idx.split("(")[1].rstrip(")")
            global_distances[source] = df_distances.loc[idx].values
            
    all_x = []
    all_y = []
    for source in df_f1.columns:
        if source in global_distances:
            x_vals = global_distances[source]
            y_vals = df_f1[source].values
            all_x.extend(x_vals)
            all_y.extend(y_vals)
    all_x_arr = np.array(all_x).reshape(-1, 1)
    all_y_arr = np.array(all_y)
    
    model_global = LinearRegression().fit(all_x_arr, all_y_arr)
    r2 = model_global.score(all_x_arr, all_y_arr)
    slope = model_global.coef_[0]
    intercept = model_global.intercept_
    
    scatter_fig = figure(width=600, height=600, tools="pan,wheel_zoom,reset,save", y_range=(0, 1), 
                         title="Scatter Plot: Distance vs F1")
    source_colors = {
        "es-digital-paragraph-degradation-seq": "blue",
        "es-digital-line-degradation-seq": "green",
        "es-digital-seq": "red",
        "es-digital-zoom-degradation-seq": "orange",
        "es-digital-rotation-degradation-seq": "purple",
        "es-digital-rotation-zoom-degradation-seq": "brown",
        "es-render-seq": "cyan"
    }
    for source in df_f1.columns:
        if source in global_distances:
            x_vals = global_distances[source]
            y_vals = df_f1[source].values
            data = {"x": x_vals, "y": y_vals, "Fuente": [source]*len(x_vals)}
            cds = ColumnDataSource(data=data)
            scatter_fig.circle('x', 'y', size=8, alpha=0.7, source=cds,
                               fill_color=source_colors.get(source, "gray"),
                               line_color=source_colors.get(source, "gray"),
                               legend_label=source)
    scatter_fig.xaxis.axis_label = "Distance (Global, por Colegio)"
    scatter_fig.yaxis.axis_label = "F1 Score"
    scatter_fig.legend.location = "top_right"
    hover_tool = HoverTool(tooltips=[("Distance", "@x"), ("F1", "@y"), ("Subset", "@Fuente")])
    scatter_fig.add_tools(hover_tool)
    # scatter_fig.match_aspect = True
    
    x_line = np.linspace(all_x_arr.min(), all_x_arr.max(), 100)
    y_line = model_global.predict(x_line.reshape(-1, 1))
    scatter_fig.line(x_line, y_line, line_width=2, line_color="black", legend_label="Global Regression")
    
    results = {
        "R2": r2,
        "slope": slope,
        "intercept": intercept,
        "scatter_fig": scatter_fig,
        "dfs_reduced": dfs_reduced,
        "unique_subsets": unique_subsets,
        "df_distances": df_distances,
        "explained_variance": explained_variance,
        "trustworthiness": trust,
        "continuity": cont
    }
    
    if reduction_method == "PCA":
        results["pca_model"] = reducer  # Agregamos el objeto PCA para usarlo luego en los plots
    
    return results

def optimize_tsne_params(df_combined, embedding_cols, df_f1, distance_metric):
    perplexity_range = np.linspace(30, 50, 10)
    learning_rate_range = np.linspace(200, 1000, 20)
    
    best_R2 = -np.inf
    best_params = None
    total_steps = len(perplexity_range) * len(learning_rate_range)
    step = 0

    progress_text = st.empty()
    
    for p in perplexity_range:
        for lr in learning_rate_range:
            step += 1
            progress_text.text(f"Evaluating: Perplexity={p:.2f}, Learning Rate={lr:.2f} (Step {step}/{total_steps})")
            
            tsne_params = {"perplexity": p, "learning_rate": lr}
            result = compute_global_regression(df_combined, embedding_cols, tsne_params, df_f1, reduction_method="t-SNE", distance_metric=distance_metric)
            r2_temp = result["R2"]
            st.write(f"Parameters: Perplexity={p:.2f}, Learning Rate={lr:.2f} -> R²={r2_temp:.4f}")
            
            if r2_temp > best_R2:
                best_R2 = r2_temp
                best_params = (p, lr)
    
    progress_text.text("Optimization completed!")
    return best_params, best_R2

def run_model(model_name):
    version = st.selectbox("Select Model Version:", options=["vanilla", "finetuned_real"], key=f"version_{model_name}")
    # Selector para el método de cómputo del embedding
    embedding_computation = st.selectbox("¿Cómo se computa el embedding?", options=["averaged", "weighted"], key=f"embedding_method_{model_name}")
    # Se asigna el prefijo correspondiente

    if embedding_computation == "weighted":
         selected_weight_factor = st.selectbox(
             "Seleccione el Weight Factor", 
             options=[0.05, 0.1, 0.25, 0.5], 
             index=0,  # índice 1 para que por defecto sea 0.05
             key=f"weight_factor_{model_name}"
         )
         weight_factor = f"{selected_weight_factor}_"
    else:
         weight_factor = ""
    
    embeddings = load_embeddings(model_name, version, embedding_computation, weight_factor)
    if embeddings is None:
        return

    # Nuevo selector para incluir o excluir el dataset pretrained
    include_pretrained = st.checkbox("Incluir dataset pretrained", value=False, key=f"legend_{model_name}_pretrained")
    if not include_pretrained:
        # Removemos la entrada pretrained del diccionario, si existe.
        embeddings.pop("pretrained", None)
    
    # Extraer columnas de embedding de los datos "real"
    embedding_cols = [col for col in embeddings["real"].columns if col.startswith("dim_")]
    # Concatenamos los datasets disponibles (ahora, sin pretrained si se deseleccionó)
    df_combined = pd.concat(list(embeddings.values()), ignore_index=True)
    
    try:
        df_f1 = pd.read_csv("data/f1-donut.csv", sep=';', index_col=0)
    except Exception as e:
        st.error(f"Error loading f1-donut.csv: {e}")
        return

    st.markdown('<h6 class="sub-title">Select Dimensionality Reduction Method</h6>', unsafe_allow_html=True)
    reduction_method = st.selectbox("", options=["PCA", "t-SNE"], key=f"reduction_{model_name}")
    
    distance_metric = st.selectbox("Select Distance Metric:", 
                                   options=["Euclidean", "Wasserstein", "KL"], 
                                   key=f"distance_metric_{model_name}")
    
    tsne_params = {}
    if reduction_method == "t-SNE":
        if st.button("Optimize TSNE parameters", key=f"optimize_tsne_{model_name}"):
            st.info("Running optimization, this can take a while...")
            best_params, best_R2 = optimize_tsne_params(df_combined, embedding_cols, df_f1, distance_metric.lower())
            st.success(f"Best parameters: Perplexity = {best_params[0]:.2f}, Learning Rate = {best_params[1]:.2f} with R² = {best_R2:.4f}")
            tsne_params = {"perplexity": best_params[0], "learning_rate": best_params[1]}
        else:
            perplexity_val = st.number_input(
                "Perplexity", 
                min_value=5.0, 
                max_value=50.0, 
                value=30.0, 
                step=1.0, 
                format="%.2f",
                key=f"perplexity_{model_name}"
            )
            learning_rate_val = st.number_input(
                "Learning Rate", 
                min_value=10.0, 
                max_value=1000.0, 
                value=200.0, 
                step=10.0, 
                format="%.2f",
                key=f"learning_rate_{model_name}"
            )
            tsne_params = {"perplexity": perplexity_val, "learning_rate": learning_rate_val}
    
    result = compute_global_regression(df_combined, embedding_cols, tsne_params, df_f1, reduction_method=reduction_method, distance_metric=distance_metric.lower())
    
    reg_metrics = pd.DataFrame({
        "Slope": [result["slope"]],
        "Intercept": [result["intercept"]],
        "R2": [result["R2"]]
    })
    st.table(reg_metrics)
    
    if reduction_method == "PCA" and result["explained_variance"] is not None:
        st.subheader("Explained Variance Ratio")
        component_names = [f"PC{i+1}" for i in range(len(result["explained_variance"]))]
        variance_df = pd.DataFrame({
            "Component": component_names,
            "Explained Variance": result["explained_variance"]
        })
        st.table(variance_df)
    elif reduction_method == "t-SNE":
        st.subheader("t-SNE Quality Metrics")
        st.write(f"Trustworthiness: {result['trustworthiness']:.4f}")
        st.write(f"Continuity: {result['continuity']:.4f}")

    # Mostrar los plots de loadings si se usó PCA (para el conjunto combinado)
    if reduction_method == "PCA" and result.get("pca_model") is not None:
        pca_model = result["pca_model"]
        components = pca_model.components_  # Shape: (n_components, n_features)
        
        st.subheader("Pesos de las Componentes Principales (Loadings) - Conjunto Combinado")
        for i, comp in enumerate(components):
            source = ColumnDataSource(data=dict(
                dimensions=embedding_cols,
                weight=comp
            ))
            p = figure(x_range=embedding_cols, title=f"Componente Principal {i+1}",
                       plot_height=400, plot_width=600,
                       toolbar_location="above",
                       tools="pan,wheel_zoom,reset,save,hover",
                       active_scroll="wheel_zoom")

            # Establecer fondo blanco
            p.background_fill_color = "white"
            # Mostrar solo grilla horizontal
            p.xgrid.grid_line_color = None
            p.ygrid.grid_line_color = "gray"
            p.vbar(x='dimensions', top='weight', width=0.8, source=source)
            p.xaxis.major_label_text_font_size = '0pt'
            hover = HoverTool(tooltips=[("Dimensión", "@dimensions"), ("Peso", "@weight")])
            p.add_tools(hover)
            p.xaxis.axis_label = "Dimensiones originales"
            p.yaxis.axis_label = "Peso"
            st.bokeh_chart(p)
    
    data_table, df_table, source_table = create_table(result["df_distances"])
    real_subset_names = list(df_table.columns[1:])
    real_select = Select(title="", value=real_subset_names[0], options=real_subset_names)
    reset_button = Button(label="Reset Colors", button_type="primary")
    line_source = ColumnDataSource(data={'x': [], 'y': []})
    
    if (reduction_method == "t-SNE" and N_COMPONENTS == 2) or (reduction_method == "PCA" and N_COMPONENTS == 2):
        fig, real_renderers, synthetic_renderers, pretrained_renderers = create_figure(
            result["dfs_reduced"],
            result["unique_subsets"],
            get_color_maps(result["unique_subsets"]),
            model_name
        )
        fig.line('x', 'y', source=line_source, line_width=2, line_color='black')
        centers_real = calculate_cluster_centers(result["dfs_reduced"]["real"], result["unique_subsets"]["real"])
        real_centers_js = {k: [v[0], v[1]] for k, v in centers_real.items()}
        synthetic_centers = {}
        synth_labels = sorted(result["dfs_reduced"]["synthetic"]['label'].unique().tolist())
        for label in synth_labels:
            subset = result["dfs_reduced"]["synthetic"][result["dfs_reduced"]["synthetic"]['label'] == label]
            if 'x' in subset.columns and 'y' in subset.columns:
                synthetic_centers[label] = [subset['x'].mean(), subset['y'].mean()]
        callback = CustomJS(args=dict(source=source_table, line_source=line_source,
                                  synthetic_centers=synthetic_centers,
                                  real_centers=real_centers_js,
                                  real_select=real_select),
        code="""
            var selected = source.selected.indices;
            if (selected.length > 0) {
                var idx = selected[0];
                var data = source.data;
                var synth_label = data['Synthetic'][idx];
                var real_label = real_select.value;
                var syn_coords = synthetic_centers[synth_label];
                var real_coords = real_centers[real_label];
                line_source.data = {'x': [syn_coords[0], real_coords[0]], 'y': [syn_coords[1], real_coords[1]]};
                line_source.change.emit();
            } else {
                line_source.data = {'x': [], 'y': []};
                line_source.change.emit();
            }
        """)
        source_table.selected.js_on_change('indices', callback)
        real_select.js_on_change('value', callback)
    
        reset_callback = CustomJS(args=dict(line_source=line_source),
        code="""
            line_source.data = {'x': [], 'y': []};
            line_source.change.emit();
        """)
        reset_button.js_on_event("button_click", reset_callback)
        layout = column(fig, result["scatter_fig"], column(real_select, reset_button, data_table))
    else:
        layout = column(result["scatter_fig"], column(real_select, reset_button, data_table))
    
    st.bokeh_chart(layout, use_container_width=True)

    buffer = io.BytesIO()
    df_table.to_excel(buffer, index=False)
    buffer.seek(0)
    
    st.download_button(
        label="Export Table",
        data=buffer,
        file_name=f"cluster_distances_{model_name}.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        key=f"download_button_excel_{model_name}"
    )

    # Nuevo bloque: PCA solo para df_real
    if reduction_method == "PCA":
        st.markdown("## PCA - Solo Muestras Reales")
        # Extraemos únicamente las muestras reales
        df_real_only = embeddings["real"].copy()
        pca_real = PCA(n_components=N_COMPONENTS)
        reduced_real = pca_real.fit_transform(df_real_only[embedding_cols].values)
        df_real_only['embedding'] = list(reduced_real)
        if reduced_real.shape[1] == 2:
            df_real_only['x'] = reduced_real[:, 0]
            df_real_only['y'] = reduced_real[:, 1]
        explained_variance_real = pca_real.explained_variance_ratio_
        unique_labels_real = sorted(df_real_only['label'].unique().tolist())
        
        # Definir mapeo de colores usando la paleta Reds9
        num_labels = len(unique_labels_real)
        if num_labels <= 9:
            red_palette = Reds9[:num_labels]
        else:
            red_palette = (Reds9 * ((num_labels // 9) + 1))[:num_labels]
        real_color_mapping = {label: red_palette[i] for i, label in enumerate(unique_labels_real)}
        
        st.subheader("PCA - Real: Explained Variance Ratio")
        component_names_real = [f"PC{i+1}" for i in range(len(explained_variance_real))]
        variance_df_real = pd.DataFrame({
            "Component": component_names_real,
            "Explained Variance": explained_variance_real
        })
        st.table(variance_df_real)

        # Mostrar los plots de loadings (Component Loadings)
        st.subheader("PCA - Real: Component Loadings")
        st.markdown("### Pesos de las Componentes Principales (Loadings) - Conjunto Combinado")
        for i, comp in enumerate(pca_real.components_):
            source = ColumnDataSource(data=dict(
                dimensions=embedding_cols,
                weight=comp
            ))
            p = figure(
                x_range=embedding_cols,
                title=f"Componente Principal {i+1}",
                plot_height=400,
                plot_width=600,
                toolbar_location="above",
                tools="pan,wheel_zoom,reset,save,hover",
                active_scroll="wheel_zoom"
            )
            # Fondo blanco y solo grid horizontal
            p.background_fill_color = "white"
            p.xgrid.grid_line_color = None
            p.ygrid.grid_line_color = "gray"
            p.vbar(x='dimensions', top='weight', width=0.8, source=source,
                fill_color="#2b83ba", line_color="#2b83ba")
            # No se muestran etiquetas en el eje horizontal
            p.xaxis.axis_label = "Dimensiones Originales"
            p.xaxis.major_label_text_font_size = '0pt'
            # Configurar el HoverTool
            hover = p.select_one(HoverTool)
            hover.tooltips = [("Dimensión", "@dimensions"), ("Peso", "@weight")]
            st.bokeh_chart(p)
        
        # Segundo PCA: Proyección de todos los subconjuntos usando los loadings calculados con df_real_only
        st.subheader("PCA - Todos los subconjuntos proyectados (usando loadings de df_real)")

        # Crear un diccionario para almacenar las proyecciones usando el PCA calculado con las muestras reales (pca_real)
        df_all = {}

        # Proyectar las muestras reales
        df_real_proj = embeddings["real"].copy()
        proj_real = pca_real.transform(df_real_proj[embedding_cols].values)
        df_real_proj['pc1'] = proj_real[:, 0]
        df_real_proj['pc2'] = proj_real[:, 1]
        df_all["real"] = df_real_proj

        # Proyectar el subconjunto synthetic, si existe
        if "synthetic" in embeddings:
            df_synth_proj = embeddings["synthetic"].copy()
            proj_synth = pca_real.transform(df_synth_proj[embedding_cols].values)
            df_synth_proj['pc1'] = proj_synth[:, 0]
            df_synth_proj['pc2'] = proj_synth[:, 1]
            df_all["synthetic"] = df_synth_proj

        # Proyectar el subconjunto pretrained, si existe
        if "pretrained" in embeddings:
            df_pretr_proj = embeddings["pretrained"].copy()
            proj_pretr = pca_real.transform(df_pretr_proj[embedding_cols].values)
            df_pretr_proj['pc1'] = proj_pretr[:, 0]
            df_pretr_proj['pc2'] = proj_pretr[:, 1]
            df_all["pretrained"] = df_pretr_proj

        # Para utilizar las mismas funciones de plot (create_figure, add_dataset_to_fig, add_synthetic_dataset_to_fig),
        # renombramos las columnas 'pc1' y 'pc2' a 'x' y 'y' en cada dataframe
        for key in df_all:
            df_all[key]["x"] = df_all[key]["pc1"]
            df_all[key]["y"] = df_all[key]["pc2"]

        # Construir los subconjuntos únicos con la granularidad deseada:
        # - Para "real" y "pretrained": agrupamos por label.
        # - Para "synthetic": agrupamos por la columna "source" (cada source tendrá sus labels).
        unique_subsets = {}
        # Real:
        unique_subsets["real"] = sorted(df_all["real"]['label'].unique().tolist())
        # Synthetic:
        if "synthetic" in df_all:
            unique_synth = {}
            for source in df_all["synthetic"]["source"].unique():
                unique_synth[source] = sorted(df_all["synthetic"][df_all["synthetic"]["source"] == source]['label'].unique().tolist())
            unique_subsets["synthetic"] = unique_synth
        else:
            unique_subsets["synthetic"] = {}
        # Pretrained:
        if "pretrained" in df_all:
            unique_subsets["pretrained"] = sorted(df_all["pretrained"]['label'].unique().tolist())
        else:
            unique_subsets["pretrained"] = []

        # Obtener los mapeos de colores utilizando la función ya definida
        color_maps = get_color_maps(unique_subsets)

        # Definir un mapeo de marcadores para los subconjuntos synthetic (granularidad por source)
        marker_mapping = {
            "es-digital-paragraph-degradation-seq": "x",
            "es-digital-line-degradation-seq": "cross",
            "es-digital-seq": "triangle",
            "es-digital-rotation-degradation-seq": "diamond",
            "es-digital-zoom-degradation-seq": "asterisk",
            "es-render-seq": "inverted_triangle"
        }

        # Ahora, crear la figura utilizando las funciones existentes para mantener la granularidad:
        # Se plotean las muestras reales, synthetic (por source) y pretrained con sus respectivos marcadores y colores.
        fig_all = figure(
            title="PCA - Todos los subconjuntos proyectados",
            plot_width=600,
            plot_height=600,
            tools="pan,wheel_zoom,reset,save",
            active_scroll="wheel_zoom",
            background_fill_color="white",
            tooltips=TOOLTIPS
        )
        # Solo grid horizontal
        fig_all.xgrid.grid_line_color = None
        fig_all.ygrid.grid_line_color = "gray"

        # Ploteamos los puntos de las muestras reales (agrupados por label)
        for label in unique_subsets["real"]:
            subset = df_all["real"][df_all["real"]['label'] == label]
            source = ColumnDataSource(data={
                'x': subset['x'],
                'y': subset['y'],
                'label': subset['label'],
                'img': subset['img']
            })
            # Usamos 'circle' para las reales
            fig_all.circle('x', 'y', size=10,
                        fill_color=color_maps["real"][label],
                        line_color=color_maps["real"][label],
                        legend_label=f"Real: {label}",
                        source=source)

        show_real_only = st.checkbox("Show only real samples", value=True, key=f"show_real_only_{model_name}")

        if not show_real_only:

            # Ploteamos los puntos de synthetic, diferenciando cada source con su marcador
            if unique_subsets["synthetic"]:
                for source_name, labels in unique_subsets["synthetic"].items():
                    df_source = df_all["synthetic"][df_all["synthetic"]["source"] == source_name]
                    marker = marker_mapping.get(source_name, "square")
                    # Para cada label en ese source, usamos la función auxiliar
                    renderers = add_synthetic_dataset_to_fig(fig_all, df_source, labels,
                                                            marker=marker,
                                                            color_mapping=color_maps["synthetic"][source_name],
                                                            group_label=source_name)
            # Ploteamos los puntos de pretrained (agrupados por label)
            if unique_subsets["pretrained"]:
                for label in unique_subsets["pretrained"]:
                    subset = df_all["pretrained"][df_all["pretrained"]['label'] == label]
                    source = ColumnDataSource(data={
                        'x': subset['x'],
                        'y': subset['y'],
                        'label': subset['label'],
                        'img': subset['img']
                    })
                    # Usamos 'triangle' para pretrained (por ejemplo)
                    fig_all.triangle('x', 'y', size=10,
                                    fill_color=color_maps["pretrained"][label],
                                    line_color=color_maps["pretrained"][label],
                                    legend_label=f"Pretrained: {label}",
                                    source=source)

        # Calcular el centroide y el radio (usando solo las muestras reales)
        center_x = df_all["real"]['x'].mean()
        center_y = df_all["real"]['y'].mean()
        distances = np.sqrt((df_all["real"]['x'] - center_x)**2 + (df_all["real"]['y'] - center_y)**2)
        radius = distances.max()

        # Dibujar el centroide y la circunferencia en el plot
        centroid_glyph = fig_all.circle(
            x=center_x, y=center_y, size=15,
            fill_color="white", line_color="black",
            legend_label="Centroide",
            name="centroid"  # Asigna un nombre único
        )

        circumference_glyph = fig_all.circle(
            x=center_x, y=center_y, radius=radius,
            fill_color=None, line_color="black",
            line_dash="dashed",
            legend_label="Circunferencia",
            name="circumference"  # Asigna un nombre único
        )

        fig_all.xaxis.axis_label = "PC1"
        fig_all.yaxis.axis_label = "PC2"
        hover_all = fig_all.select_one(HoverTool)
        hover_all.renderers = [r for r in fig_all.renderers if r.name not in ["centroid", "circumference"]]

        # hover_all.tooltips = [("Label", "@label"), ("PC1", "@x"), ("PC2", "@y")]

        # Agregar checkbox para mostrar u ocultar la leyenda, igual que en el primer PCA
        show_legend_second = st.checkbox("Show Legend", value=False, key=f"legend_second_{model_name}")
        fig_all.legend.visible = show_legend_second
        fig_all.legend.location = "top_right"
        fig_all.match_aspect = True

        st.bokeh_chart(fig_all)

        # Mostrar el valor del radio debajo del gráfico
        st.write(f"El radio de la circunferencia (calculado a partir de las muestras reales) es: {radius:.4f}")


        # --- Cálculo de distancias y scatter plot de Distance vs F1 para el nuevo PCA ---

        # Se calcula la distancia de cada subset synthetic a cada subset real usando los datos proyectados (df_all)
        # Se utiliza la función compute_cluster_distances_synthetic_individual ya definida
        real_labels_new = sorted(df_all["real"]['label'].unique().tolist())
        df_distances_new = compute_cluster_distances_synthetic_individual(
            df_all["synthetic"],
            df_all["real"],
            real_labels_new,
            metric="wasserstein",  # Puedes cambiar la métrica según lo requieras
            bins=20
        )

        # Extraer las distancias globales (por cada source) del dataframe obtenido,
        # buscando filas cuyo índice comience con "Global" (formato "Global (source)")
        global_distances_new = {}
        for idx in df_distances_new.index:
            if idx.startswith("Global"):
                source_name = idx.split("(")[1].rstrip(")")
                global_distances_new[source_name] = df_distances_new.loc[idx].values

        # Ahora, relacionar estas distancias con los valores de F1 (ya cargados en df_f1)
        all_x_new = []
        all_y_new = []
        for source in df_f1.columns:
            if source in global_distances_new:
                x_vals = global_distances_new[source]
                y_vals = df_f1[source].values
                all_x_new.extend(x_vals)
                all_y_new.extend(y_vals)
        all_x_arr_new = np.array(all_x_new).reshape(-1, 1)
        all_y_arr_new = np.array(all_y_new)

        # Realizar la regresión lineal global sobre estos datos
        model_global_new = LinearRegression().fit(all_x_arr_new, all_y_arr_new)
        r2_new = model_global_new.score(all_x_arr_new, all_y_arr_new)
        slope_new = model_global_new.coef_[0]
        intercept_new = model_global_new.intercept_

        # Crear el scatter plot
        scatter_fig_new = figure(
            width=600,
            height=600,
            tools="pan,wheel_zoom,reset,save,hover",
            active_scroll="wheel_zoom",
            title="Scatter Plot: Distance vs F1 (Nueva PCA)",
            background_fill_color="white",
            y_range=(0, 1)
        )
        # Configurar únicamente grid horizontal
        scatter_fig_new.xgrid.grid_line_color = None
        scatter_fig_new.ygrid.grid_line_color = "gray"
        scatter_fig_new.match_aspect = True

        # Mantenemos el mismo código de colores que en el otro scatter plot
        source_colors = {
            "es-digital-paragraph-degradation-seq": "blue",
            "es-digital-line-degradation-seq": "green",
            "es-digital-seq": "red",
            "es-digital-zoom-degradation-seq": "orange",
            "es-digital-rotation-degradation-seq": "purple",
            "es-digital-rotation-zoom-degradation-seq": "brown",
            "es-render-seq": "cyan"
        }

        # Dibujar cada conjunto: para cada source (por ejemplo, es-render-seq, etc.)
        for source in df_f1.columns:
            if source in global_distances_new:
                x_vals = global_distances_new[source]
                y_vals = df_f1[source].values
                data = {"x": x_vals, "y": y_vals, "Fuente": [source]*len(x_vals)}
                cds = ColumnDataSource(data=data)
                scatter_fig_new.circle(
                    'x', 'y', size=8, alpha=0.7, source=cds,
                    fill_color=source_colors.get(source, "gray"),
                    line_color=source_colors.get(source, "gray"),
                    legend_label=source
                )

        scatter_fig_new.xaxis.axis_label = "Distance (Global, por Colegio) - Nueva PCA"
        scatter_fig_new.yaxis.axis_label = "F1 Score"
        scatter_fig_new.legend.location = "top_right"

        hover_tool_new = scatter_fig_new.select_one(HoverTool)
        hover_tool_new.tooltips = [("Distance", "@x"), ("F1", "@y"), ("Subset", "@Fuente")]

        # Dibujar la línea de regresión global
        x_line_new = np.linspace(all_x_arr_new.min(), all_x_arr_new.max(), 100)
        y_line_new = model_global_new.predict(x_line_new.reshape(-1,1))
        scatter_fig_new.line(x_line_new, y_line_new, line_width=2, line_color="black", legend_label="Global Regression")

        st.bokeh_chart(scatter_fig_new)

        st.write(f"Regresión global (Nueva PCA): R² = {r2_new:.4f}, Slope = {slope_new:.4f}, Intercept = {intercept_new:.4f}")

        # --- INICIO DEL BLOQUE: Heatmap de características ---
        st.markdown("## Heatmap de Características")

        try:
            df_heat = pd.read_csv("data/heatmaps.csv")
            # Si fuera necesario, se pueden limpiar los nombres de las columnas:
            # df_heat.columns = [col.strip("'\"") for col in df_heat.columns]
        except Exception as e:
            st.error(f"Error al cargar heatmaps.csv: {e}")
            df_heat = None

        if df_heat is not None:
            # Verificamos que la columna 'img' esté presente en df_all["real"]
            if 'img' not in df_all["real"].columns:
                st.error("La columna 'img' no se encuentra en las muestras reales para hacer el merge con heatmaps.csv.")
            else:
                # Crear la columna 'name' extrayendo el nombre final de la URL y removiendo ".png"
                df_all["real"]["name"] = df_all["real"]["img"].apply(
                    lambda x: x.split("/")[-1].replace(".png", "") if isinstance(x, str) else x
                )
                
                # Hacemos merge de las posiciones reales con el CSV de heatmaps usando la columna 'name'
                df_heatmap = pd.merge(df_all["real"], df_heat, on="name", how="inner")
                
                # Extraemos las características disponibles (excluyendo 'name')
                feature_options = [col for col in df_heat.columns if col != "name"]
                selected_feature = st.selectbox("Seleccione la característica para el heatmap:", 
                                                options=feature_options, key=f"heatmap_{model_name}")
                
                # Determinar el rango de las posiciones (x, y) de las muestras reales
                x_min, x_max = df_heatmap['x'].min(), df_heatmap['x'].max()
                y_min, y_max = df_heatmap['y'].min(), df_heatmap['y'].max()
                
                # Definir resolución de la rejilla (por ejemplo, 50x50)
                grid_size = 50
                x_bins = np.linspace(x_min, x_max, grid_size + 1)
                y_bins = np.linspace(y_min, y_max, grid_size + 1)
                
                # Si la variable seleccionada no es numérica, la convertimos a códigos numéricos
                # y guardamos la correspondencia para la leyenda.
                cat_mapping = None
                if df_heatmap[selected_feature].dtype == bool or not pd.api.types.is_numeric_dtype(df_heatmap[selected_feature]):
                    cat = df_heatmap[selected_feature].astype('category')
                    cat_mapping = list(cat.cat.categories)
                    df_heatmap[selected_feature] = cat.cat.codes
                
                # Intentamos calcular el heatmap; si falla, aplicamos la conversión a categoría
                try:
                    heat_stat, x_edges, y_edges, binnumber = binned_statistic_2d(
                        df_heatmap['x'], df_heatmap['y'], df_heatmap[selected_feature],
                        statistic='mean', bins=[x_bins, y_bins]
                    )
                except TypeError:
                    cat = df_heatmap[selected_feature].astype('category')
                    cat_mapping = list(cat.cat.categories)
                    df_heatmap[selected_feature] = cat.cat.codes
                    heat_stat, x_edges, y_edges, binnumber = binned_statistic_2d(
                        df_heatmap['x'], df_heatmap['y'], df_heatmap[selected_feature],
                        statistic='mean', bins=[x_bins, y_bins]
                    )
                
                # La función image de Bokeh espera una lista de arrays; se transpone para alinear los ejes.
                heatmap_data = heat_stat.T
                
                # Crear el mapa de color
                color_mapper = LinearColorMapper(palette="Viridis256", low=np.nanmin(heatmap_data), high=np.nanmax(heatmap_data), nan_color = 'rgba(0, 0, 0, 0)')
                
                # Crear la figura para el heatmap con fondo blanco
                heatmap_fig = figure(title=f"Heatmap de '{selected_feature}'",
                                    x_range=(x_min, x_max), y_range=(y_min, y_max),
                                    width=600, height=600,
                                    tools="pan,wheel_zoom,reset,save", active_scroll="wheel_zoom", tooltips=TOOLTIPS)
                
                # Dibujar el heatmap usando la imagen
                heatmap_fig.image(image=[heatmap_data], x=x_min, y=y_min,
                                dw=x_max - x_min, dh=y_max - y_min,
                                color_mapper=color_mapper)
                
                # Crear la barra de colores
                color_bar = ColorBar(color_mapper=color_mapper, location=(0, 0))
                # Si se usó conversión a categoría, formateamos la barra para mostrar las etiquetas originales
                if cat_mapping is not None:
                    
                    # Creamos ticks fijos solo para cada categoría
                    ticks = list(range(len(cat_mapping)))
                    color_bar.ticker = FixedTicker(ticks=ticks)
                    
                    categories_json = json.dumps(cat_mapping)
                    color_bar.formatter = FuncTickFormatter(code=f"""
                        var categories = {categories_json};
                        var index = Math.round(tick);
                        if(index >= 0 && index < categories.length) {{
                            return categories[index];
                        }} else {{
                            return "";
                        }}
                    """)
                heatmap_fig.add_layout(color_bar, 'right')



                # Agregar renderer de puntos invisibles para tooltips
                source_points = ColumnDataSource(data={
                    'x': df_heatmap['x'],
                    'y': df_heatmap['y'],
                    'img': df_heatmap['img'],
                    'label': df_heatmap['name']  # Asegúrate de que esta columna exista; si no, usa otra
                })
                # Dibujar círculos con transparencia total (no se verán)
                invisible_renderer = heatmap_fig.circle('x', 'y', size=10, source=source_points, fill_alpha=0, line_alpha=0.5)

                hover_tool_points = HoverTool(renderers=[invisible_renderer], tooltips=TOOLTIPS)
                heatmap_fig.add_tools(hover_tool_points)
                
                st.bokeh_chart(heatmap_fig)

def main():
    config_style()
    tabs = st.tabs(["Donut", "Idefics2"])
    with tabs[0]:
        st.markdown('<h2 class="sub-title">Donut 🤗</h2>', unsafe_allow_html=True)
        run_model("Donut")
    with tabs[1]:
        st.markdown('<h2 class="sub-title">Idefics2 🤗</h2>', unsafe_allow_html=True)
        run_model("Idefics2")

if __name__ == "__main__":
    main()