import gradio as gr
from gradio_leaderboard import Leaderboard, ColumnFilter, SelectColumns
import pandas as pd

# Static data
STATIC_DATA = [
    ["VLM", "w/o WM", "–", "RGB", "72B", 50.27, 6.24],
    ["Image Gen.", "PathDreamer [36]", "Viewpoint", "RGB-D; Sem; Pano", "0.69B", 56.99, 5.28],
    ["Image Gen.", "SE3DS [11]", "Viewpoint", "RGB-D; Pano", "1.1B", 57.53, 5.29],
    ["Video Gen.", "NWM [25]", "Trajectory", "RGB", "1B", 57.35, 5.68],
    ["Video Gen.", "SVD [6]", "Image", "RGB", "1.5B", 57.71, 5.29],
    ["Video Gen.", "LTX-Video [5]", "Text", "RGB", "2B", 56.08, 5.37],
    ["Video Gen.", "Hunyuan [4]", "Text", "RGB", "13B", 57.71, 5.21],
    ["Video Gen.", "Wan2.1 [23]", "Text", "RGB", "14B", 58.26, 5.24],
    ["Video Gen.", "Cosmos [1]", "Text", "RGB", "2B", 52.27, 5.898],
    ["Video Gen.", "Runway", "Text", "–", "–", "–", "–"],
    ["Video Gen. Post-Train", "SVD† [6]", "Action", "RGB; Pano", "1.5B", 60.98, 5.02],
    ["Video Gen. Post-Train", "LTX† [5]", "Action", "RGB; Pano", "2B", 57.53, 5.49],
    ["Video Gen. Post-Train", "WAN2.1† [23]", "Action", "RGB; Pano", "14B", "XXX", "XXX"],
    ["Video Gen. Post-Train", "Cosmos† [1]", "Action", "RGB; Pano", "2B", 60.25, 5.08],
]

COLUMNS = ["Model Type", "Method", "Control Type", "Input Type", "#Param.", "Acc. ↑", "Mean Traj. ↓"]
LEADERBOARD_DF = pd.DataFrame(STATIC_DATA, columns=COLUMNS)

# Custom CSS (simplified)
custom_css = """
/* Add any custom styling here */
.gradio-container {
    max-width: 1200px !important;
}
"""

def init_leaderboard(dataframe):
    if dataframe is None or dataframe.empty:
        raise ValueError("Leaderboard DataFrame is empty or None.")
    return Leaderboard(
        value=dataframe,
        datatype=["str", "str", "str", "str", "str", "number", "number"],
        select_columns=SelectColumns(
            default_selection=COLUMNS,
            cant_deselect=["Model Type", "Method", "Acc. ↑"],
            label="Select Columns to Display:",
        ),
        search_columns=["Model Type", "Method"],
        hide_columns=[],
        filter_columns=[
            ColumnFilter("Model Type", type="checkboxgroup", label="Model types"),
            ColumnFilter("Control Type", type="checkboxgroup", label="Control types"),
            ColumnFilter("Input Type", type="checkboxgroup", label="Input types"),
        ],
        bool_checkboxgroup_label="Hide models",
        interactive=False,
    )

demo = gr.Blocks(css=custom_css, title="Model Performance Leaderboard")
with demo:
    gr.HTML("<h1 style='text-align: center'>🏆 Model Performance Leaderboard</h1>")
    gr.Markdown("""
    **Performance comparison across vision-language models, image generation, and video generation models.**
    
    📊 **Metrics:** Acc. ↑ (Accuracy - higher is better) | Mean Traj. ↓ (Mean Trajectory error - lower is better)
    """, elem_classes="markdown-text")

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🏅 Leaderboard", elem_id="leaderboard-tab", id=0):
            leaderboard = init_leaderboard(LEADERBOARD_DF)

        with gr.TabItem("📝 About", elem_id="about-tab", id=1):
            gr.Markdown("""
            # About This Leaderboard
            
            This leaderboard showcases performance metrics across different types of AI models:
            
            ## Model Categories
            - **VLM**: Vision-Language Models
            - **Image Gen.**: Image Generation Models  
            - **Video Gen.**: Video Generation Models
            - **Video Gen. Post-Train**: Post-training specialized Video Generation Models
            
            ## Metrics Explained
            - **Acc. ↑**: Accuracy score (higher values indicate better performance)
            - **Mean Traj. ↓**: Mean trajectory error (lower values indicate better performance)
            
            ## Notes
            - † indicates post-training specialized models
            - XXX indicates results pending/unavailable
            - – indicates not applicable or not available
            
            *Results may vary across different evaluation settings and benchmarks.*
            """, elem_classes="markdown-text")

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