challenges/answer_therapy/page.py

import json
import os
from datetime import datetime, timezone

import gradio as gr
import pandas as pd

from auth_utils import get_user_greeting, get_daily_cap_info
from hf_utils import (
    _load_leaderboard_df,
    _load_user_submissions,
    _create_submission_record,
    _count_submissions_today,
    _get_cap_for_phase,
    _today_utc_str,
)
from config import DAILY_SUBMISSION_CAP, SUBMISSIONS_TOKEN
from challenges.answer_therapy.config import (
    PHASES, LEADERBOARD_METRICS, DEFAULT_SORT_METRIC,
    LEADERBOARD_FILE, CHALLENGE_PHASE, SUBFOLDER, CHALLENGE_TYPE,
    EVAL_DETAILS_MD, FORMAT_MD,
)
from challenges.answer_therapy.validate import validate_submission


def load_leaderboard():
    """Load VQA Answer Therapy leaderboard, filtered to challenge phase."""
    try:
        df = _load_leaderboard_df(LEADERBOARD_FILE, LEADERBOARD_METRICS)
    except Exception as e:
        return pd.DataFrame(), f"❌ Could not load leaderboard: {e}"

    if not df.empty and "phase_codename" in df.columns:
        df = df[df["phase_codename"] == CHALLENGE_PHASE]

    if df.empty:
        return pd.DataFrame(), "ℹ️ No scored Challenge phase submissions yet. Be the first!"

    df = df.sort_values(by=DEFAULT_SORT_METRIC, ascending=False, kind="mergesort")
    df_display = df.copy()
    df_display.insert(0, "Rank", range(1, len(df_display) + 1))
    if "timestamp" in df_display.columns:
        df_display["Scored At"] = pd.to_datetime(
            df_display["timestamp"], unit="s", errors="coerce"
        ).dt.strftime("%d %b %Y, %I:%M %p")
        df_display.drop(columns=["timestamp"], inplace=True)
    for col in ["username", "email", "phase_codename", "submission_id"]:
        if col in df_display.columns:
            df_display.drop(columns=[col], inplace=True)

    rename = {
        "team": "Team", "model": "Model",
        "overall_f1":        "Overall F1",
        "overall_precision": "Precision (Overall)",
        "overall_recall":    "Recall (Overall)",
        "vizwiz_f1":         "VizWiz F1",
        "vizwiz_precision":  "Precision (VizWiz)",
        "vizwiz_recall":     "Recall (VizWiz)",
        "vqav2_f1":          "VQAv2 F1",
        "vqa_precision":     "Precision (VQA)",
        "vqa_recall":        "Recall (VQA)",
    }
    df_display.rename(columns=rename, inplace=True)
    return df_display, ""


def handle_submit(file, team, model_name, phase_label, profile: gr.OAuthProfile | None):
    """Handle VQA Answer Therapy submission."""
    if profile is None:
        return "❌ You must be logged in with your HuggingFace account to submit.", ""
    username = profile.username
    email = getattr(profile, "email", "") or ""

    if not SUBMISSIONS_TOKEN:
        return "❌ Missing SUBMISSIONS_TOKEN. Add it in Space Settings → Secrets.", ""
    if file is None:
        return "❌ Please upload a JSON file.", ""
    if not team.strip():
        return "❌ Please enter a Team / Display Name.", ""
    if not model_name.strip():
        return "❌ Please enter a Model Name.", ""

    phase_codename = next((p["codename"] for p in PHASES if p["label"] == phase_label), phase_label)
    cap = _get_cap_for_phase(phase_codename)
    subs_today = _count_submissions_today(username, phase_codename)
    if subs_today >= cap:
        phase_str = "challenge" if "challenge" in phase_codename else "this"
        return f"⛔ You've reached your daily limit of {cap} submission(s) for the {phase_str} phase. Come back tomorrow!", ""

    try:
        with open(file, "r", encoding="utf-8") as f:
            pred_obj = json.load(f)
    except Exception:
        return "❌ Could not parse JSON file.", ""

    ok, msg = validate_submission(pred_obj)
    if not ok:
        return f"❌ Invalid submission format: {msg}", ""

    original_filename = os.path.basename(file)
    try:
        submission_id = _create_submission_record(
            pred=pred_obj,
            team=team,
            model_name=model_name,
            phase_codename=phase_codename,
            challenge_type=CHALLENGE_TYPE,
            original_filename=original_filename,
            username=username,
            email=email,
            subfolder=SUBFOLDER,
        )
    except Exception as e:
        return f"❌ Upload failed: {e}", ""

    remaining = cap - subs_today - 1
    return (
        f"✅ Submission queued! Visit **My Submissions** to track results. "
        f"You have {remaining}/{cap} submissions remaining today for this phase.",
        submission_id,
    )


def load_my_submissions(phase_filter: str, profile: gr.OAuthProfile | None):
    if profile is None:
        return pd.DataFrame(), "❌ Please log in to view your submissions.", ""
    username = profile.username
    submissions = _load_user_submissions(username, subfolder=SUBFOLDER)

    if not submissions:
        return pd.DataFrame(), "ℹ️ No submissions yet. Head to Submit Predictions to get started!", ""

    all_submissions = submissions[:]
    if phase_filter and phase_filter != "All":
        submissions = [s for s in submissions if s["phase"] == phase_filter]
    if not submissions:
        return pd.DataFrame(), f"ℹ️ No submissions found for phase **{phase_filter}**.", ""

    state_icons = {"queued": "🟡", "running": "🔵", "done": "🟢", "failed": "🔴", "unknown": "⚪"}
    df = pd.DataFrame(submissions)
    if "timestamp" in df.columns:
        df["Submitted At"] = pd.to_datetime(
            df["timestamp"], unit="s", errors="coerce"
        ).dt.strftime("%d %b %Y, %I:%M %p")
    if "state" in df.columns:
        df["Status"] = df["state"].apply(lambda s: f"{state_icons.get(s, '⚪')} {s.capitalize()}")

    display_cols = ["Submitted At", "Status", "team", "model", "phase", "error"]
    metric_cols  = [m for m in LEADERBOARD_METRICS if m in df.columns]
    display_cols += metric_cols
    df_display = df[[c for c in display_cols if c in df.columns]].copy()
    df_display.rename(columns={
        "team": "Team", "model": "Model", "phase": "Phase", "error": "Error",
        "overall_f1": "Overall F1", "overall_precision": "Precision", "overall_recall": "Recall",
        "vizwiz_f1": "VizWiz F1", "vqav2_f1": "VQAv2 F1",
    }, inplace=True)
    for m in metric_cols:
        if m in df_display.columns:
            df_display[m] = df_display[m].apply(lambda x: f"{x:.2f}" if pd.notna(x) else "")

    total = len(all_submissions)
    done  = sum(1 for s in all_submissions if s["state"] == "done")
    today_count = sum(
        1 for s in all_submissions
        if datetime.fromtimestamp(s["timestamp"], tz=timezone.utc).strftime("%Y-%m-%d") == _today_utc_str()
    )
    stats = (
        f"**Total:** {total}  |  "
        f"**Scored:** {done}  |  "
        f"**Today:** {today_count}/{DAILY_SUBMISSION_CAP}"
    )
    return df_display, "", stats


def build_at_page(demo: gr.Blocks) -> None:
    with demo.route("VQA Answer Therapy", "/answer-therapy") as at_page:

        with gr.Row():
            with gr.Column(scale=5):
                gr.Markdown("# 📍 Answer Therapy Challenge")
                gr.Markdown(
                    "Predict whether answers to a visual question all share the same image region. "
                    "Evaluated with F1, Precision, and Recall across VizWiz and VQAv2 question sets."
                )
            with gr.Column(scale=1, min_width=160):
                gr.LoginButton(size="lg")
                at_greeting = gr.Markdown("👋 Log in to submit.")

        gr.Markdown("---")

        with gr.Tabs():

            # ── Submit ──
            with gr.TabItem("🚀 Submit Predictions"):
                with gr.Row():
                    at_submit_greeting = gr.Markdown("👋 Log in with HuggingFace to submit.")
                    at_cap_info        = gr.Markdown("")
                gr.Markdown("---")
                with gr.Row():
                    with gr.Column(scale=3):
                        gr.Markdown("#### Upload Submission File")
                        at_file_input = gr.File(label="Choose a JSON file", file_types=[".json"])
                        with gr.Accordion("📄 Submission Format", open=False):
                            gr.Markdown(FORMAT_MD)
                    with gr.Column(scale=2):
                        gr.Markdown("#### Submission Info")
                        at_team_input  = gr.Textbox(label="Team / Display Name", placeholder="e.g. My Awesome Team")
                        at_model_input = gr.Textbox(label="Model Name", placeholder="e.g. CLIP-ViT-L")
                        at_phase_input = gr.Dropdown(
                            label="Phase",
                            choices=[p["label"] for p in PHASES],
                            value=PHASES[0]["label"],
                        )
                at_submit_btn    = gr.Button("Submit (Queue for Evaluation)", variant="primary", size="lg")
                at_submit_status = gr.Markdown("")
                at_sid_box       = gr.Code(label="Submission ID", language=None, visible=False)

                def do_at_submit(file, team, model_name, phase_label, profile: gr.OAuthProfile | None):
                    msg, sid = handle_submit(file, team, model_name, phase_label, profile)
                    return msg, gr.update(value=sid, visible=bool(sid))

                at_submit_btn.click(
                    do_at_submit,
                    inputs=[at_file_input, at_team_input, at_model_input, at_phase_input],
                    outputs=[at_submit_status, at_sid_box],
                )

                def update_at_submit_ui(profile: gr.OAuthProfile | None):
                    return get_user_greeting(profile), get_daily_cap_info(profile, PHASES)

                at_page.load(update_at_submit_ui, outputs=[at_submit_greeting, at_cap_info])

            # ── My Submissions ──
            with gr.TabItem("📋 My Submissions"):
                at_my_sub_greeting = gr.Markdown("👋 Log in with HuggingFace to view your submissions.")
                at_my_sub_stats    = gr.Markdown("")
                with gr.Row():
                    at_phase_filter = gr.Dropdown(
                        label="Filter by Phase",
                        choices=["All"] + [p["codename"] for p in PHASES],
                        value="All",
                        scale=2,
                    )
                    at_refresh_my_btn = gr.Button("🔄 Refresh", variant="secondary", scale=1)
                at_my_sub_msg   = gr.Markdown("")
                at_my_sub_table = gr.Dataframe(interactive=False, wrap=True)

                def refresh_at_my_subs(phase_filter, profile: gr.OAuthProfile | None):
                    df, msg, stats = load_my_submissions(phase_filter, profile)
                    return df, msg, stats, get_user_greeting(profile)

                at_refresh_my_btn.click(
                    refresh_at_my_subs,
                    inputs=[at_phase_filter],
                    outputs=[at_my_sub_table, at_my_sub_msg, at_my_sub_stats, at_my_sub_greeting],
                )
                at_phase_filter.change(
                    refresh_at_my_subs,
                    inputs=[at_phase_filter],
                    outputs=[at_my_sub_table, at_my_sub_msg, at_my_sub_stats, at_my_sub_greeting],
                )
                at_page.load(
                    refresh_at_my_subs,
                    inputs=[at_phase_filter],
                    outputs=[at_my_sub_table, at_my_sub_msg, at_my_sub_stats, at_my_sub_greeting],
                )

            # ── Leaderboard ──
            with gr.TabItem("🏆 Leaderboard"):
                gr.Markdown("### Challenge Phase Rankings")
                gr.Markdown("Ranked by **Overall F1** (descending). Challenge phase only.")
                with gr.Accordion("📐 How is the Score Calculated?", open=False):
                    gr.Markdown(EVAL_DETAILS_MD)
                at_lb_msg   = gr.Markdown("")
                at_lb_table = gr.Dataframe(interactive=False, wrap=True)
                at_refresh_lb_btn = gr.Button("🔄 Refresh Leaderboard", variant="secondary", size="sm")

                def refresh_at_leaderboard(profile: gr.OAuthProfile | None):
                    df, msg = load_leaderboard()
                    return df, msg, get_user_greeting(profile)

                at_refresh_lb_btn.click(refresh_at_leaderboard, outputs=[at_lb_table, at_lb_msg, at_greeting])
                at_page.load(refresh_at_leaderboard, outputs=[at_lb_table, at_lb_msg, at_greeting])