import gradio as gr
import plotly.graph_objects as go
import numpy as np


from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets


def create_heatmap(selected_models, benchmark):
    if not selected_models:
        return gr.update(visible=False)
    
    # Generate random similarity matrix
    size = len(selected_models)
    similarities = np.random.rand(size, size)
    
    # Create symmetric matrix (for demo purposes)
    similarities = (similarities + similarities.T) / 2
    
    # Create heatmap with Plotly
    fig = go.Figure(data=go.Heatmap(
        z=similarities,
        x=selected_models,
        y=selected_models,
        colorscale='Viridis',
        hoverongaps=False
    ))
    
    fig.update_layout(
        title=f"Model Similarity for {benchmark}",
        xaxis_title="Models",
        yaxis_title="Models",
        height=600,
        width=800
    )
    
    return fig

with gr.Blocks(title="LLM Similarity Analyzer") as demo:
    gr.Markdown("## Model Similarity Comparison Tool")
    
    # Model selection section
    with gr.Row():
        dataset_dropdown = gr.Dropdown(
                choices=get_leaderboard_datasets(),
                label="Select Dataset",
                filterable=True,
                interactive=True,
                info="Leaderboard benchmark datasets"
            )

        model_dropdown = gr.Dropdown(
            choices=get_leaderboard_models_cached(),
            label="Select Models",
            multiselect=True,
            filterable=True,
            allow_custom_value=False,
            info="Search and select multiple models (click selected models to remove)"
        )
    
    # Heatmap display
    heatmap = gr.Plot(
        label="Similarity Heatmap",
        visible=False,
        container=False
    )
    
    # Interactive updates
    model_dropdown.input(
        fn=create_heatmap,
        inputs=(model_dropdown, dataset_dropdown),
        outputs=heatmap
    )


if __name__ == "__main__":
    demo.launch()