import gradio as gr
import pandas as pd
import os

# Load data from CSV files
DATA_DIR = "data"

def load_csv_data(filename):
    """Load data from CSV file"""
    filepath = os.path.join(DATA_DIR, filename)
    if os.path.exists(filepath):
        return pd.read_csv(filepath)
    else:
        return pd.DataFrame()

# Load datasets
CLASSIFICATION_DF = load_csv_data("classification.csv")
DETECTION_DF = load_csv_data("detection.csv")
SEGMENTATION_DF = load_csv_data("segmentation.csv")


def filter_and_search(df, search, datasets, models, organizations, columns):
    filtered = df.copy()

    if search:
        mask = filtered.apply(lambda row: row.astype(str).str.contains(search, case=False).any(), axis=1)
        filtered = filtered[mask]

    if datasets:
        filtered = filtered[filtered["Dataset"].isin(datasets)]
    if models:
        filtered = filtered[filtered["Model"].isin(models)]
    if organizations:
        filtered = filtered[filtered["Organization"].isin(organizations)]

    if columns:
        display_cols = [col for col in columns if col in filtered.columns]
        filtered = filtered[display_cols]

    return filtered


def build_tab(df, name):
    """Build a leaderboard tab from a DataFrame"""
    if df.empty:
        return gr.TabItem(name)

    # Pivot the dataframe to have columns like "Dataset1 (Metric1)", "Dataset1 (Metric2)", etc.
    datasets = sorted(df["Dataset"].unique().tolist())
    models = sorted(df["Model"].unique().tolist())
    organizations = sorted(df["Organization"].unique().tolist())
    metric_cols = [col for col in df.columns if col not in ["Model", "Organization", "Dataset", "Author", "Author Link"]]

    # Create pivoted dataframe
    pivoted_data = []
    for model in models:
        for org in organizations:
            model_org_data = df[(df["Model"] == model) & (df["Organization"] == org)]
            if not model_org_data.empty:
                # Get Author and Author Link from the first entry (they should be the same for model+org)
                author = model_org_data["Author"].values[0]
                author_link = model_org_data["Author Link"].values[0]

                # Format as markdown link if author_link exists
                if pd.notna(author_link) and author_link.strip():
                    author_display = f"[{author}]({author_link})"
                else:
                    author_display = author

                row = {"Model": model, "Organization": org, "Author": author_display}
                for dataset in datasets:
                    dataset_data = model_org_data[model_org_data["Dataset"] == dataset]
                    if not dataset_data.empty:
                        for metric in metric_cols:
                            col_name = f"{dataset} ({metric})"
                            row[col_name] = dataset_data[metric].values[0]
                    else:
                        for metric in metric_cols:
                            col_name = f"{dataset} ({metric})"
                            row[col_name] = "-"
                pivoted_data.append(row)

    pivoted_df = pd.DataFrame(pivoted_data)

    # Build column list for selector
    metric_combo_cols = []
    for dataset in datasets:
        for metric in metric_cols:
            metric_combo_cols.append(f"{dataset} ({metric})")

    all_cols = ["Model", "Organization", "Author"] + metric_combo_cols

    with gr.TabItem(name, elem_id="llm-benchmark-tab-table"):
        with gr.Row():
            with gr.Column(scale=4):
                search_bar = gr.Textbox(
                    label="Search",
                    placeholder="Separate multiple queries with ';'",
                    elem_id="search-bar"
                )

                # Column selector for base columns only (not dataset+metric combos)
                base_cols = ["Model", "Organization", "Author"]
                col_selector = gr.CheckboxGroup(
                    choices=base_cols,
                    value=base_cols,
                    label="Select Columns to Display:",
                    elem_classes="column-select"
                )

                # Set datatype to 'markdown' for Author column to enable clickable links
                datatypes = []
                for col in pivoted_df.columns:
                    if col == "Author":
                        datatypes.append("markdown")
                    else:
                        datatypes.append("str")

                table = gr.Dataframe(
                    value=pivoted_df,
                    elem_id="leaderboard-table",
                    interactive=False,
                    wrap=True,
                    datatype=datatypes
                )

            with gr.Column(scale=1):
                gr.Markdown("**Model types**")
                model_filter = gr.CheckboxGroup(
                    choices=models,
                    value=models,
                    label="",
                    elem_classes="filter-group"
                )

                gr.Markdown("**Organizations**")
                org_filter = gr.CheckboxGroup(
                    choices=organizations,
                    value=organizations,
                    label="",
                    elem_classes="filter-group"
                )

                gr.Markdown("**Datasets**")
                dataset_filter = gr.CheckboxGroup(
                    choices=datasets,
                    value=datasets,
                    label="",
                    elem_classes="filter-group"
                )

                gr.Markdown("**Metrics**")
                metric_filter = gr.CheckboxGroup(
                    choices=metric_cols,
                    value=metric_cols,
                    label="",
                    elem_classes="filter-group"
                )

        def update(search, md, org, dset, metrics, cols):
            filtered = pivoted_df.copy()

            if search:
                mask = filtered.apply(lambda row: row.astype(str).str.contains(search, case=False).any(), axis=1)
                filtered = filtered[mask]

            if md:
                filtered = filtered[filtered["Model"].isin(md)]
            if org:
                filtered = filtered[filtered["Organization"].isin(org)]

            # Build display columns based on selected base columns and dataset/metric filters
            display_cols = []

            # Add selected base columns
            for col in cols:
                if col in base_cols:
                    display_cols.append(col)

            # Add metric columns that match selected datasets and metrics
            for col in metric_combo_cols:
                # Check if this column matches selected datasets and metrics
                col_dataset = col.split(" (")[0]
                col_metric = col.split(" (")[1].rstrip(")")
                if col_dataset in dset and col_metric in metrics:
                    display_cols.append(col)

            filtered = filtered[display_cols]

            return filtered

        search_bar.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)
        model_filter.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)
        org_filter.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)
        dataset_filter.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)
        metric_filter.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)
        col_selector.change(update, [search_bar, model_filter, org_filter, dataset_filter, metric_filter, col_selector], table)


custom_css = """
.markdown-text {
    font-size: 16px !important;
}

#leaderboard-table {
    margin-top: 15px;
}

#leaderboard-table table {
    width: 100%;
    table-layout: auto;
}

#leaderboard-table thead th {
    font-weight: bold;
    text-align: center;
    padding: 12px 8px;
    white-space: normal;
    word-wrap: break-word;
}

#leaderboard-table tbody td {
    text-align: center;
    padding: 10px 8px;
    white-space: nowrap;
}

#leaderboard-table tbody td:first-child {
    text-align: left;
    font-weight: 500;
    min-width: 120px;
    max-width: 200px;
    white-space: nowrap;
}

#leaderboard-table thead th:first-child {
    text-align: left;
    font-weight: bold;
    min-width: 120px;
    max-width: 200px;
    white-space: nowrap;
}

#leaderboard-table tbody td:nth-child(2),
#leaderboard-table thead th:nth-child(2) {
    text-align: left;
    min-width: 100px;
}

#leaderboard-table tbody td:nth-child(3),
#leaderboard-table thead th:nth-child(3) {
    text-align: left;
    min-width: 120px;
}

/* Style links in Author column */
#leaderboard-table a {
    color: #0066cc;
    text-decoration: none;
}

#leaderboard-table a:hover {
    text-decoration: underline;
}

#search-bar {
    margin-bottom: 0px;
}

#search-bar textarea {
    border: 1px solid #e0e0e0 !important;
    border-radius: 8px !important;
}

.tab-buttons button {
    font-size: 20px;
}

.filter-group {
    margin-bottom: 1em;
}

.filter-group label {
    font-size: 14px;
}

.column-select {
    margin-bottom: 1.5em;
}

.column-select label {
    display: flex;
    flex-wrap: wrap;
    gap: 0.5em;
}

.column-select label > span {
    display: inline-flex;
    align-items: center;
}
"""

TITLE = """<h1 align="center" id="space-title">Mars-Bench Leaderboard</h1>"""

INTRO = """
A comprehensive benchmark for evaluating computer vision models on Mars-specific datasets.
This leaderboard tracks model performance across three key tasks: classification, segmentation, and object detection.
"""

demo = gr.Blocks(css=custom_css, title="Mars-Bench Leaderboard")
with demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRO, elem_classes="markdown-text")

    with gr.Tabs(elem_classes="tab-buttons"):
        build_tab(CLASSIFICATION_DF, "🏅 Classification")
        build_tab(SEGMENTATION_DF, "🏅 Segmentation")
        build_tab(DETECTION_DF, "🏅 Object Detection")

        with gr.TabItem("📝 Submit", elem_id="submit-tab"):
            gr.Markdown("""
            # Submit Your Model Results

            To submit your model's results to Mars-Bench, please provide the following information.
            All submissions will be reviewed before being added to the leaderboard.
            """)

            gr.Markdown("""
            ### How to submit:
            1. Fill out the form below
            2. Click "Generate Submission Text"
            3. Copy the generated text
            4. Go to the [Community tab](https://huggingface.co/spaces/gremlin97/MarsBoard/discussions)
            5. Click "New discussion"
            6. Paste the text and submit

            ---
            """)

            with gr.Row():
                with gr.Column():
                    submit_task = gr.Dropdown(
                        choices=["Classification", "Segmentation", "Object Detection"],
                        label="Task Type",
                        info="Select the task category"
                    )
                    submit_model = gr.Textbox(
                        label="Model Name",
                        placeholder="e.g., ResNet-50",
                        info="Name of your model"
                    )
                    submit_org = gr.Textbox(
                        label="Organization",
                        placeholder="e.g., Microsoft, Google",
                        info="Your organization or affiliation"
                    )
                    submit_dataset = gr.Textbox(
                        label="Dataset",
                        placeholder="e.g., DoMars16, Mars Crater",
                        info="Dataset used for evaluation"
                    )

                with gr.Column():
                    submit_author = gr.Textbox(
                        label="Author",
                        placeholder="e.g., K. He",
                        info="Name of the author"
                    )
                    submit_author_link = gr.Textbox(
                        label="Author Page Link (Optional)",
                        placeholder="https://scholar.google.com/... or https://github.com/...",
                        info="Link to author's Google Scholar, GitHub, or personal page"
                    )
                    submit_metrics = gr.Textbox(
                        label="Metrics (JSON format)",
                        placeholder='{"Accuracy": 95.8, "F1-Score": 94.9}',
                        info="Provide metrics in JSON format",
                        lines=3
                    )
                    submit_paper = gr.Textbox(
                        label="Paper Link (Optional)",
                        placeholder="https://arxiv.org/abs/...",
                        info="Link to your research paper"
                    )
                    submit_code = gr.Textbox(
                        label="Code Repository (Optional)",
                        placeholder="https://github.com/...",
                        info="Link to your code repository"
                    )
                    submit_email = gr.Textbox(
                        label="Contact Email",
                        placeholder="your.email@example.com",
                        info="We'll contact you about your submission"
                    )

            submit_notes = gr.Textbox(
                label="Additional Notes (Optional)",
                placeholder="Training details, hyperparameters, reproduction instructions...",
                lines=4,
                info="Any additional information"
            )

            generate_btn = gr.Button("Generate Submission Text", variant="primary", size="lg")
            submission_output = gr.Textbox(
                label="Copy this text and create a new discussion in the Community tab",
                lines=15,
                interactive=True
            )

            gr.Markdown("""
            We'll review your submission and add it to the leaderboard if approved.
            """)

            def generate_submission_text(task, model, org, dataset, author, author_link, metrics, paper, code, email, notes):
                if not all([task, model, org, dataset, author, metrics, email]):
                    return "❌ Error: Please fill in all required fields (Task Type, Model Name, Organization, Dataset, Author, Metrics, Contact Email)"

                submission_text = f"""## New Model Submission

**Task Type:** {task}
**Model Name:** {model}
**Organization:** {org}
**Dataset:** {dataset}
**Author:** {author}
**Author Page Link:** {author_link if author_link else "N/A"}

### Metrics
```json
{metrics}
```

### Links
- **Paper:** {paper if paper else "N/A"}
- **Code:** {code if code else "N/A"}

### Contact
**Email:** {email}

### Additional Notes
{notes if notes else "N/A"}

---
*Please review this submission for inclusion in Mars-Bench.*
"""
                return submission_text

            generate_btn.click(
                generate_submission_text,
                inputs=[submit_task, submit_model, submit_org, submit_dataset, submit_author,
                       submit_author_link, submit_metrics, submit_paper, submit_code, submit_email, submit_notes],
                outputs=submission_output
            )

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