app.py

import os

import gradio as gr
import pandas as pd
from apscheduler.schedulers.background import BackgroundScheduler
from gradio_leaderboard import Leaderboard, SelectColumns
from huggingface_hub import snapshot_download

from src.about import (
    INTRODUCTION_TEXT,
    LLM_BENCHMARKS_TEXT,
    TITLE,
)
from src.display.css_html_js import custom_css
from src.display.formatting import styled_error, styled_message
from src.display.utils import (
    COLS,
    SUBMISSION_COLS,
    SUBMISSION_TYPES,
    TeamColumn,
    fields,
)
from src.envs import (
    API,
    REPO_ID,
    SUBMISSIONS_PATH,
    SUBMISSIONS_REPO,
    TEAMS_PATH,
    TEAMS_REPO,
    TOKEN,
)
from src.evaluation.load_labels import load_true_labels
from src.populate import get_leaderboard_df, get_submission_queue_df
from src.submission.submit_csv import submit_csv
from src.teams.register import create_team


def restart_space():
    API.restart_space(repo_id=REPO_ID)


try:
    snapshot_download(
        repo_id=TEAMS_REPO,
        local_dir=TEAMS_PATH,
        repo_type="dataset",
        tqdm_class=None,
        etag_timeout=30,
        token=TOKEN,
    )
except Exception as e:
    print(f"Warning: Could not download teams dataset from {TEAMS_REPO}: {e}")

try:
    snapshot_download(
        repo_id=SUBMISSIONS_REPO,
        local_dir=SUBMISSIONS_PATH,
        repo_type="dataset",
        tqdm_class=None,
        etag_timeout=30,
        token=TOKEN,
    )
except Exception as e:
    print(f"Warning: Could not download submissions dataset from {SUBMISSIONS_REPO}: {e}")

try:
    load_true_labels()
except Exception as e:
    print(f"Warning: Could not load true labels: {e}")

LEADERBOARD_DF = get_leaderboard_df(SUBMISSIONS_PATH, COLS)

(
    accepted_submissions_df,
    rejected_submissions_df,
    all_submissions_df,
) = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)


def init_leaderboard(dataframe):
    team_columns = [c for c in fields(TeamColumn) if isinstance(c, type(TeamColumn.team_name))]

    valid_cols = [col for col in COLS if col is not None and isinstance(col, str) and col.strip() != ""]
    if not valid_cols:
        valid_cols = [
            "Team Name",
            "Best F1 Score ⬆️",
            "Best Accuracy ⬆️",
            "Best Precision ⬆️",
            "Best Recall ⬆️",
            "Best TP ⬆️",
            "Best FP ⬇️",
            "Best FN ⬇️",
            "Best TN ⬆️",
            "Submission Date",
        ]

    if dataframe is None or dataframe.empty:
        empty_df = pd.DataFrame(columns=valid_cols)
        column_to_type = {c.name: c.type for c in team_columns}
        datatypes = []
        for col in valid_cols:
            dtype = column_to_type.get(col, "str")
            if not dtype or dtype == "":
                dtype = "str"
            datatypes.append(dtype)

        print(empty_df)
        print(datatypes)

        search_col = TeamColumn.team_name.name if TeamColumn.team_name.name in valid_cols else valid_cols[0]

        return Leaderboard(
            value=empty_df,
            datatype=datatypes,
            search_columns=[search_col],
            select_columns=SelectColumns(
                default_selection=valid_cols,
                cant_deselect=[search_col],
                label="Select Columns to Display:",
            ),
            filter_columns=[],
            hide_columns=[],
            interactive=False,
            height=800,
        )

    dataframe = dataframe[
        [col for col in dataframe.columns if col is not None and isinstance(col, str) and col.strip() != ""]
    ]

    if dataframe.empty or len(dataframe.columns) == 0:
        dataframe = pd.DataFrame(columns=valid_cols)

    missing_cols = [col for col in valid_cols if col not in dataframe.columns]
    for col in missing_cols:
        dataframe[col] = None

    dataframe = dataframe[valid_cols]

    column_to_type = {c.name: c.type for c in team_columns}
    datatypes = []
    for col in dataframe.columns:
        dtype = column_to_type.get(col, "str")
        if not dtype or dtype == "":
            dtype = "str"
        datatypes.append(dtype)

    default_selection = [
        c.name for c in team_columns if getattr(c, "displayed_by_default", False) and c.name in dataframe.columns
    ]

    cant_deselect = [c.name for c in team_columns if getattr(c, "never_hidden", False) and c.name in dataframe.columns]

    hide_cols = [c.name for c in team_columns if getattr(c, "hidden", False) and c.name in dataframe.columns]

    search_cols = []
    if hasattr(TeamColumn, "team_name") and hasattr(TeamColumn.team_name, "name"):
        search_col_name = TeamColumn.team_name.name
        if search_col_name and search_col_name in dataframe.columns:
            search_cols = [search_col_name]

    return Leaderboard(
        value=dataframe,
        datatype=datatypes,
        select_columns=SelectColumns(
            default_selection=default_selection if default_selection else [],
            cant_deselect=cant_deselect if cant_deselect else [],
            label="Select Columns to Display:",
        ),
        search_columns=search_cols if search_cols else [],
        hide_columns=hide_cols if hide_cols else [],
        filter_columns=[],
        interactive=False,
        height=800,
    )


def register_team_ui(team_name: str, num_teammates: int):
    try:
        num_teammates_int = int(num_teammates)
    except (ValueError, TypeError):
        return styled_error("Number of teammates must be a valid integer.")

    try:
        token, team_data = create_team(team_name, num_teammates_int)
        return styled_message(
            f"Team '{team_name}' registered successfully!\n\n"
            f"**IMPORTANT: Save your token now - you won't be able to see it again!**\n\n"
            f"Your team token: `{token}`\n\n"
            f"Use this token to submit your predictions."
        )
    except ValueError as e:
        return styled_error(str(e))
    except Exception as e:
        return styled_error(f"Registration failed: {str(e)}")


def submit_csv_ui(token: str, csv_file):
    updated_leaderboard_df = get_leaderboard_df(SUBMISSIONS_PATH, COLS)

    if not token or not token.strip():
        accepted, rejected, all_subs = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)
        return styled_error("Please provide your team token."), updated_leaderboard_df, accepted, rejected, all_subs

    if csv_file is None:
        accepted, rejected, all_subs = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)
        return styled_error("Please upload a CSV file."), updated_leaderboard_df, accepted, rejected, all_subs

    try:
        with open(csv_file.name, "r") as f:
            csv_content = f.read()
    except Exception as e:
        accepted, rejected, all_subs = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)
        return styled_error(f"Could not read CSV file: {str(e)}"), updated_leaderboard_df, accepted, rejected, all_subs

    success, message = submit_csv(token, csv_content)

    updated_leaderboard_df = get_leaderboard_df(SUBMISSIONS_PATH, COLS)
    accepted, rejected, all_subs = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)

    if success:
        return styled_message(message), updated_leaderboard_df, accepted, rejected, all_subs
    else:
        return styled_error(message), updated_leaderboard_df, accepted, rejected, all_subs


def refresh_leaderboard():
    fresh_df = get_leaderboard_df(SUBMISSIONS_PATH, COLS)
    return fresh_df


def refresh_submission_history():
    accepted, rejected, all_subs = get_submission_queue_df(SUBMISSIONS_PATH, SUBMISSION_COLS)
    return accepted, rejected, all_subs


logo_image_path = os.path.abspath("assets/logo.png")

demo = gr.Blocks(css=custom_css, theme=gr.themes.Soft())
with demo:
    with gr.Row(elem_id="title-row"):
        with gr.Column(scale=0, min_width=120):
            gr.Image(
                value=logo_image_path,
                show_label=False,
                container=False,
                height=120,
                width=120,
                show_download_button=False,
                show_fullscreen_button=False,
                show_share_button=False,
                interactive=False,
            )
        with gr.Column(scale=1):
            gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, 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(LLM_BENCHMARKS_TEXT, elem_classes="markdown-text")

        with gr.TabItem("👥 Register Team", elem_id="register-tab", id=2):
            with gr.Column():
                gr.Markdown("## Create Your Team", elem_classes="markdown-text")
                gr.Markdown(
                    "Register your team to participate in the hackathon. "
                    "You will receive a token that you'll need to submit predictions.",
                    elem_classes="markdown-text",
                )

                with gr.Row():
                    with gr.Column():
                        team_name_input = gr.Textbox(
                            label="Team Name",
                            placeholder="Enter your team name",
                            interactive=True,
                        )
                        num_teammates_input = gr.Number(
                            label="Number of Teammates",
                            value=1,
                            minimum=1,
                            maximum=100,
                            step=1,
                            interactive=True,
                        )
                        register_button = gr.Button("Register Team", variant="primary")
                        registration_result = gr.Markdown()

                register_button.click(
                    register_team_ui,
                    [team_name_input, num_teammates_input],
                    registration_result,
                )

        with gr.TabItem("🚀 Submit Predictions", elem_id="submit-tab", id=3):
            with gr.Column():
                gr.Markdown("## Submit Your Predictions", elem_classes="markdown-text")
                gr.Markdown(
                    "Upload a CSV file with your predictions. The CSV must have two columns: "
                    "`id` (UUID) and `label` (must be exactly `0.0` or `1.0`). All IDs from the test set must be included.",
                    elem_classes="markdown-text",
                )

                with gr.Row():
                    with gr.Column():
                        token_input = gr.Textbox(
                            label="Team Token",
                            placeholder="Enter your team token",
                            type="password",
                            interactive=True,
                        )
                        csv_file_input = gr.File(
                            label="CSV File",
                            file_types=[".csv"],
                            interactive=True,
                        )
                        submit_button = gr.Button("Submit CSV", variant="primary")
                        submission_result = gr.Markdown()

                with gr.Accordion("📊 Submission History", open=False):
                    with gr.Tabs():
                        with gr.TabItem("✅ Accepted Submissions"):
                            accepted_table = gr.components.Dataframe(
                                value=accepted_submissions_df,
                                datatype=SUBMISSION_TYPES,
                                row_count=10,
                            )
                        with gr.TabItem("❌ Rejected Submissions"):
                            rejected_table = gr.components.Dataframe(
                                value=rejected_submissions_df,
                                datatype=SUBMISSION_TYPES,
                                row_count=10,
                            )
                        with gr.TabItem("📋 All Submissions"):
                            all_table = gr.components.Dataframe(
                                value=all_submissions_df,
                                datatype=SUBMISSION_TYPES,
                                row_count=10,
                            )

                submit_button.click(
                    submit_csv_ui,
                    [token_input, csv_file_input],
                    [submission_result, leaderboard, accepted_table, rejected_table, all_table],
                )

    demo.load(
        lambda: (refresh_leaderboard(), *refresh_submission_history()),
        outputs=[leaderboard, accepted_table, rejected_table, all_table],
    )

if __name__ == "__main__":
    scheduler = BackgroundScheduler()
    scheduler.add_job(restart_space, "interval", seconds=1800)
    scheduler.start()
    demo.queue(default_concurrency_limit=40).launch(
        allowed_paths=[logo_image_path],
    )