import matplotlib
matplotlib.use('Agg')
import gradio as gr


from ui_components import (
    create_leaderboard_display,
    get_full_leaderboard_data,
    create_winners_by_category_html,
)

from content import (
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
    INTRO_PARAGRAPH
)

from visualizations import (
    create_evolution_over_time_chart,
    create_accuracy_by_size_chart
)

from constants import MARK_BY_DEFAULT

# --- Global State for Viewers (simple caching) ---
CACHED_VIEWERS = {}
CACHED_TAG_MAPS = {}

def build_page():
    with gr.Row(elem_id="intro-row"):
        with gr.Column(scale=1):
            gr.HTML(INTRO_PARAGRAPH, elem_id="intro-paragraph")

    # --- Leaderboard Display Section ---
    gr.Markdown("---")
    CATEGORY_NAME = "Overall"
    gr.HTML(f'<h2>OpenHands Index {CATEGORY_NAME} Leaderboard <span style="font-weight: normal; color: inherit;">(Aggregate)</span></h2>', elem_id="main-header")

    test_df, test_tag_map = get_full_leaderboard_data("test")
    if not test_df.empty:
        # Get the checkbox and dropdown returned from create_leaderboard_display
        show_open_only_checkbox, mark_by_dropdown = create_leaderboard_display(
            full_df=test_df,
            tag_map=test_tag_map,
            category_name=CATEGORY_NAME,
            split_name="test"
        )
        
        # Prepare open-only filtered dataframe for Winners and Evolution
        if 'Openness' in test_df.columns:
            test_df_open = test_df[test_df['Openness'].str.lower() == 'open'].copy()
        else:
            test_df_open = test_df.copy()
        
        # --- Winners by Category Section ---
        gr.Markdown("---")
        gr.HTML('<h2>Winners by Category</h2>', elem_id="winners-header")
        gr.Markdown("Top 5 performing systems in each benchmark category.")
        
        # Create both all and open-only versions of winners HTML
        winners_html_all = create_winners_by_category_html(test_df, top_n=5)
        winners_html_open = create_winners_by_category_html(test_df_open, top_n=5)
        
        winners_component = gr.HTML(winners_html_all, elem_id="winners-by-category")
        
        # --- New Visualization Sections ---
        gr.Markdown("---")
        
        # Evolution Over Time Section
        gr.HTML('<h2>Evolution Over Time</h2>', elem_id="evolution-header")
        gr.Markdown("Track how model performance has improved over time based on release dates.")
        
        # Create initial evolution chart with default mark_by
        evolution_fig_all = create_evolution_over_time_chart(test_df, MARK_BY_DEFAULT)
        
        evolution_component = gr.Plot(value=evolution_fig_all, elem_id="evolution-chart")
        
        gr.Markdown("---")
        
        # Open Model Accuracy by Size Section (always shows open models only by design)
        gr.HTML('<h2>Open Model Accuracy by Size</h2>', elem_id="size-accuracy-header")
        gr.Markdown("Compare open-weights model performance against their parameter count.")
        
        size_fig = create_accuracy_by_size_chart(test_df, MARK_BY_DEFAULT)
        size_component = gr.Plot(value=size_fig, elem_id="size-accuracy-chart")
        
        # Update function for Winners, Evolution, and Size charts based on filters
        def update_extra_sections(show_open_only, mark_by):
            # Select the appropriate dataframe based on open_only filter
            df_to_use = test_df_open if show_open_only else test_df
            
            # Winners HTML (not affected by mark_by, only open_only)
            winners_html = winners_html_open if show_open_only else winners_html_all
            
            # Regenerate charts with current mark_by setting
            evolution_fig = create_evolution_over_time_chart(df_to_use, mark_by)
            size_fig = create_accuracy_by_size_chart(test_df, mark_by)  # Size chart always uses full df (filters internally)
            
            return winners_html, evolution_fig, size_fig
        
        # Connect both checkbox and dropdown to update all extra sections
        if show_open_only_checkbox is not None:
            show_open_only_checkbox.change(
                fn=update_extra_sections,
                inputs=[show_open_only_checkbox, mark_by_dropdown],
                outputs=[winners_component, evolution_component, size_component]
            )
        
        if mark_by_dropdown is not None:
            mark_by_dropdown.change(
                fn=update_extra_sections,
                inputs=[show_open_only_checkbox if show_open_only_checkbox else gr.State(value=False), mark_by_dropdown],
                outputs=[winners_component, evolution_component, size_component]
            )
        
    else:
        gr.Markdown("No data available.")

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