"""Main Gradio application for the Mars Vision Leaderboard."""

import gradio as gr
from .leaderboard import update_leaderboard
from .data import TASK_DATA
import pandas as pd

TASKS = ["Classification", "Object Detection", "Segmentation"]

def get_available_models(task: str) -> list:
    """Get list of available models for a task."""
    data, _ = TASK_DATA[task]
    return sorted(list(set(data['Model'])))

def create_interface():
    """Create the Gradio interface."""
    with gr.Blocks(title="Mars Vision Leaderboard", theme=gr.themes.Default()) as demo:
        gr.Markdown(
            """
            # Mars Vision Leaderboard
            A comprehensive evaluation of computer vision models on Mars-specific datasets.
            This leaderboard tracks performance across multiple tasks including classification,
            object detection, and segmentation.
            """
        )

        with gr.Tabs():
            with gr.TabItem("📊 General Datasets"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("""
                        #### Surface Analysis
                        • **Robins & Hynek** - Craters
                        • **Lagain** - Surface Features
                        • **SPOC** - Surface Properties
                        • **AI4MARS** - Surface Analysis
                        • **MarsData** - General Surface
                        """)
                    with gr.Column(scale=1):
                        gr.Markdown("""
                        #### Classification
                        • **DoMars16k** - Surface Types
                        • **Mars Image** - Content Analysis
                        • **Deep Mars** - Deep Learning
                        • **Dusty vs Non-dusty** - Dust Analysis
                        """)
                    with gr.Column(scale=1):
                        gr.Markdown("""
                        #### Segmentation
                        • **S5Mars** - Surface
                        • **Mars-Seg** - Features
                        • **Martian Landslide**
                        • **Martian Frost**
                        """)
            
            with gr.TabItem("🎯 Specialized Tasks"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("""
                        #### Detection Tasks
                        • **Change Detection**
                        • **Outlier Detection**
                        • **Novelty Detection**
                        """)
                    with gr.Column(scale=1):
                        gr.Markdown("""
                        #### Feature Analysis
                        • **Cone Detection**
                        • **Dust Devil Tracks**
                        • **Cone Segmentation**
                        """)

        with gr.Row():
            task_dropdown = gr.Dropdown(
                choices=TASKS,
                value=TASKS[0],
                label="Select Task",
            )
            model_multiselect = gr.Dropdown(
                choices=get_available_models(TASKS[0]),
                value=None,
                label="Filter Models (Optional)",
                multiselect=True,
            )

        with gr.Column():
            gr.Markdown("### Best Performing Models Across Datasets")
            best_models_output = gr.Dataframe(
                interactive=False,
                wrap=True,
                headers=["Metric", "Rank", "Model", "Average Score"],
            )

            gr.Markdown("### Detailed Results")
            table_output = gr.Dataframe(interactive=False, wrap=True)
            
            with gr.Row():
                plot_output1 = gr.Plot(label="Performance Plot 1")
                plot_output2 = gr.Plot(label="Performance Plot 2")

        def update_models(task):
            return gr.Dropdown(choices=get_available_models(task))

        def update_with_filters(task, models):
            return update_leaderboard(task, models)

        # Event handlers
        task_dropdown.change(
            fn=update_models,
            inputs=[task_dropdown],
            outputs=[model_multiselect],
        )

        for component in [task_dropdown, model_multiselect]:
            component.change(
                fn=update_with_filters,
                inputs=[task_dropdown, model_multiselect],
                outputs=[table_output, plot_output1, plot_output2, best_models_output],
            )

        # Initial update
        demo.load(
            fn=update_with_filters,
            inputs=[task_dropdown, model_multiselect],
            outputs=[table_output, plot_output1, plot_output2, best_models_output],
        )

    return demo

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