main.py

import os
import json
import uuid
import time
import tempfile
from datetime import datetime, timezone
from typing import Any, Dict, List, Tuple, Optional

import gradio as gr
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError

# =========================
# CONFIG
# =========================

DB_REPO_ID = os.getenv("DB_REPO_ID", "VizWiz-Challenges/submissions-db")
DB_REPO_TYPE = "dataset"
SUBMISSIONS_TOKEN = os.getenv("SUBMISSIONS_TOKEN", "")

DAILY_SUBMISSION_CAP = 5

PHASES = [
    {"label": "Dev (test-dev2024)",            "codename": "test-dev2024"},
    {"label": "Standard (test-standard2024)",  "codename": "test-standard2024"},
    {"label": "Challenge (test-challenge2024)", "codename": "test-challenge2024"},
]

CHALLENGE_TYPES = ["Object Detection", "Instance Segmentation"]
LEADERBOARD_METRICS = ["bbox_mAP", "bbox_AP50", "segm_mAP", "segm_AP50"]
DEFAULT_SORT_METRIC = "segm_AP50"

# =========================
# HF API
# =========================

_api = None
def api_client() -> HfApi:
    global _api
    if _api is None:
        _api = HfApi()
    return _api

# =========================
# SUBMISSION HELPERS
# =========================

def _today_utc_str() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")

def _count_submissions_today(username: str) -> int:
    try:
        files = api_client().list_repo_files(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
        )
        today = _today_utc_str()
        count = 0
        for f in files:
            if not (f.startswith("submissions/") and f.endswith("/meta.json")):
                continue
            try:
                p = hf_hub_download(
                    repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                    filename=f, token=SUBMISSIONS_TOKEN
                )
                meta = json.load(open(p))
                if meta.get("username", "").lower() == username.lower():
                    ts = meta.get("timestamp", 0)
                    sub_date = datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%d")
                    if sub_date == today:
                        count += 1
            except Exception:
                continue
        return count
    except Exception:
        return 0

def _validate_submission_json(obj: Any) -> Tuple[bool, str]:
    if not isinstance(obj, list):
        return False, "Submission must be a JSON list of annotations."
    required_keys = {"image_id", "score", "category_id", "area", "bbox", "segmentation"}
    for i, ann in enumerate(obj):
        if not isinstance(ann, dict):
            return False, f"Annotation at index {i} must be an object/dict."
        missing = required_keys - set(ann.keys())
        if missing:
            return False, f"Annotation at index {i} missing keys: {sorted(list(missing))}"
        if not isinstance(ann["image_id"], int):
            return False, f"image_id at index {i} must be an integer."
        if not isinstance(ann["category_id"], int):
            return False, f"category_id at index {i} must be an integer."
        if not isinstance(ann["score"], (int, float)):
            return False, f"score at index {i} must be a number."
        if not isinstance(ann["area"], (int, float)):
            return False, f"area at index {i} must be a number."
        bbox = ann["bbox"]
        if not (isinstance(bbox, list) and len(bbox) == 4 and all(isinstance(x, (int, float)) for x in bbox)):
            return False, f"bbox at index {i} must be a list of 4 numbers."
        if not isinstance(ann["segmentation"], list):
            return False, f"segmentation at index {i} must be a list."
    return True, "OK"

def _upload_json(data: Any, path_in_repo: str, commit_message: str = "") -> None:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        json.dump(data, tmp, ensure_ascii=False)
        tmp_path = tmp.name
    try:
        api_client().upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=path_in_repo,
            repo_id=DB_REPO_ID,
            repo_type=DB_REPO_TYPE,
            token=SUBMISSIONS_TOKEN,
            commit_message=commit_message or f"Add {path_in_repo}",
        )
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass

def _create_submission_record(*, pred, team, model_name, phase_codename,
                               challenge_type, original_filename, username, email) -> str:
    if not SUBMISSIONS_TOKEN:
        raise ValueError("Missing SUBMISSIONS_TOKEN.")
    submission_id = str(uuid.uuid4())
    ts = int(time.time())
    meta = {
        "submission_id": submission_id,
        "team": team.strip(),
        "model": model_name.strip(),
        "phase_codename": phase_codename,
        "challenge_type": challenge_type,
        "timestamp": ts,
        "original_filename": original_filename,
        "username": username,
        "email": email,
    }
    status = {"state": "queued", "timestamp": ts}
    base = f"submissions/{submission_id}"
    _upload_json(pred, f"{base}/pred.json", f"pred {submission_id}")
    _upload_json(meta, f"{base}/meta.json", f"meta {submission_id}")
    _upload_json(status, f"{base}/status.json", f"status {submission_id}")
    return submission_id

def _load_leaderboard_df() -> pd.DataFrame:
    empty = pd.DataFrame(columns=["team", "model", "phase_codename", *LEADERBOARD_METRICS, "timestamp"])
    if not SUBMISSIONS_TOKEN:
        return empty
    try:
        path = hf_hub_download(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
            filename="leaderboard.jsonl", token=SUBMISSIONS_TOKEN
        )
    except HfHubHTTPError as e:
        if "404" in str(e):
            return empty
        raise

    rows = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue

    if not rows:
        return empty

    df = pd.DataFrame(rows)
    for col in ["team", "model", "phase_codename", "timestamp", *LEADERBOARD_METRICS]:
        if col not in df.columns:
            df[col] = None
    if DEFAULT_SORT_METRIC in df.columns:
        df = df.sort_values(by=DEFAULT_SORT_METRIC, ascending=False, kind="mergesort")
    return df

def _load_user_submissions(username: str) -> List[Dict]:
    if not SUBMISSIONS_TOKEN:
        return []
    try:
        files = api_client().list_repo_files(
            repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE, token=SUBMISSIONS_TOKEN
        )
    except Exception:
        return []

    results = []
    for f in files:
        if not (f.startswith("submissions/") and f.endswith("/meta.json")):
            continue
        try:
            sid = f.split("/")[1]
            meta_path = hf_hub_download(
                repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                filename=f, token=SUBMISSIONS_TOKEN
            )
            meta = json.load(open(meta_path))
            if meta.get("username", "").lower() != username.lower():
                continue
            try:
                status_path = hf_hub_download(
                    repo_id=DB_REPO_ID, repo_type=DB_REPO_TYPE,
                    filename=f"submissions/{sid}/status.json", token=SUBMISSIONS_TOKEN
                )
                status = json.load(open(status_path))
            except Exception:
                status = {"state": "unknown"}
            metrics = status.get("metrics", {}) if status.get("state") == "done" else {}
            error = status.get("error", "") if status.get("state") == "failed" else ""
            results.append({
                "submission_id": sid,
                "team": meta.get("team", ""),
                "model": meta.get("model", ""),
                "phase": meta.get("phase_codename", ""),
                "challenge_type": meta.get("challenge_type", ""),
                "timestamp": meta.get("timestamp", 0),
                "state": status.get("state", "unknown"),
                "error": error[:120] if error else "",
                **metrics,
            })
        except Exception:
            continue

    results.sort(key=lambda x: x["timestamp"], reverse=True)
    return results

# =========================
# GRADIO HANDLER FUNCTIONS
# =========================

def load_leaderboard():
    try:
        df = _load_leaderboard_df()
    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"] == "test-standard2024"]

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

    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)

    return df_display, ""


def handle_submit(file, team, model_name, phase_label, challenge_type, profile: gr.OAuthProfile | None):
    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.", ""

    # Daily cap check
    subs_today = _count_submissions_today(username)
    if subs_today >= DAILY_SUBMISSION_CAP:
        return f"⛔ You've reached your daily limit of {DAILY_SUBMISSION_CAP} submissions. Come back tomorrow!", ""

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

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

    # Resolve phase codename
    phase_codename = next((p["codename"] for p in PHASES if p["label"] == phase_label), phase_label)
    original_filename = os.path.basename(file)

    # Upload
    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,
        )
    except Exception as e:
        return f"❌ Upload failed: {e}", ""

    remaining = DAILY_SUBMISSION_CAP - subs_today - 1
    return (
        f"✅ Submission queued successfully! Visit **My Submissions** to see the results. You have {remaining}/{DAILY_SUBMISSION_CAP} submissions remaining today.",
    )


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)

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

    # Keep unfiltered list for accurate stats
    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", "challenge_type", "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",
        "challenge_type": "Challenge Type", "error": "Error",
    }, inplace=True)
    for m in metric_cols:
        if m in df_display.columns:
            df_display[m] = df_display[m].apply(lambda x: f"{x:.4f}" if pd.notna(x) else "")

    # Summary stats (always from unfiltered list)
    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 get_daily_cap_info(profile: gr.OAuthProfile | None):
    if profile is None:
        return ""
    subs_today = _count_submissions_today(profile.username)
    remaining = DAILY_SUBMISSION_CAP - subs_today
    return f"**{remaining}/{DAILY_SUBMISSION_CAP}** submissions remaining today"


def get_user_greeting(profile: gr.OAuthProfile | None):
    if profile is None:
        return "👋 Log in with your HuggingFace account to submit predictions."
    return f"👤 Logged in as **{profile.username}**"


# =========================
# STATIC CONTENT
# =========================

EVAL_DETAILS_MD = """
### How is the Score Calculated?

Your submission is evaluated automatically against hidden ground-truth annotations using **pycocotools**.

| Metric | Description |
|--------|-------------|
| `bbox_mAP` | Bounding box mean average precision |
| `bbox_AP50` | Bounding box AP at IoU = 0.50 |
| `segm_mAP` | Segmentation mean average precision |
| `segm_AP50` | Segmentation AP at IoU = 0.50 *(default ranking metric)* |
"""

FORMAT_MD = """
### Submission Format

Your JSON file must be a **list of annotation objects**, each containing:

```json
[
  {
    "image_id": 123,
    "category_id": 101,
    "score": 0.95,
    "area": 1024.0,
    "bbox": [x, y, width, height],
    "segmentation": [[x1, y1, x2, y2, ...]]
  },
  ...
]
```
"""

CHALLENGES = [
    {
        "id": "object-localization",
        "title": "Object Localization",
        "emoji": "🎯",
        "description": "Detect and segment objects in images taken by blind photographers. Submit bounding box and instance segmentation predictions evaluated with pycocotools.",
        "metrics": "bbox_mAP · bbox_AP50 · segm_mAP · segm_AP50",
        "route": "/object-localization",
        "active": True,
    },
    {
        "id": "vqa",
        "title": "Visual Question Answering",
        "emoji": "🤔",
        "description": "Answer open-ended questions about images taken by blind users. Models are evaluated on answer accuracy and relevance.",
        "metrics": "Coming soon",
        "route": "/vqa",
        "active": False,
    },
    {
        "id": "answer-grounding",
        "title": "Answer Grounding",
        "emoji": "📍",
        "description": "Ground free-form answers to visual regions in images taken by blind photographers.",
        "metrics": "Coming soon",
        "route": "/answer-grounding",
        "active": False,
    },
]


def _challenge_cards_html() -> str:
    cards = ""
    for c in CHALLENGES:
        if c["active"]:
            card = f"""
            <div style="border:2px solid #2563eb;border-radius:12px;padding:24px;
                        background:#f0f7ff;flex:1;min-width:220px;max-width:340px;">
                <div style="font-size:2.2rem;margin-bottom:8px;">{c['emoji']}</div>
                <h3 style="margin:0 0 8px 0;color:#1e40af;font-size:1.05rem;">{c['title']}</h3>
                <p style="margin:0 0 12px 0;color:#374151;font-size:0.88rem;line-height:1.5;">{c['description']}</p>
                <div style="background:#dbeafe;border-radius:6px;padding:5px 10px;
                            font-size:0.75rem;color:#1d4ed8;font-family:monospace;margin-bottom:12px;">
                    📊 {c['metrics']}
                </div>
                <a href="{c['route']}" style="background:#2563eb;color:white;padding:7px 16px;
                   border-radius:6px;font-size:0.85rem;font-weight:600;text-decoration:none;">
                    Enter Challenge →
                </a>
            </div>"""
        else:
            card = f"""
            <div style="border:2px solid #e5e7eb;border-radius:12px;padding:24px;
                        background:#f9fafb;flex:1;min-width:220px;max-width:340px;opacity:0.6;">
                <div style="font-size:2.2rem;margin-bottom:8px;">{c['emoji']}</div>
                <h3 style="margin:0 0 8px 0;color:#6b7280;font-size:1.05rem;">{c['title']}</h3>
                <p style="margin:0 0 12px 0;color:#9ca3af;font-size:0.88rem;line-height:1.5;">{c['description']}</p>
                <div style="background:#f3f4f6;border-radius:6px;padding:5px 10px;
                            font-size:0.75rem;color:#9ca3af;font-family:monospace;margin-bottom:12px;">
                    📊 {c['metrics']}
                </div>
                <span style="background:#e5e7eb;color:#6b7280;padding:7px 16px;
                      border-radius:6px;font-size:0.85rem;font-weight:600;">
                    🔒 Coming Soon
                </span>
            </div>"""
        cards += card
    return f'<div style="display:flex;gap:20px;flex-wrap:wrap;margin:16px 0 24px 0;">{cards}</div>'


# =========================
# BUILD UI
# =========================

with gr.Blocks() as demo:

    # ── HOME PAGE ─────────────────────────────────────────────────────────
    with gr.Row():
        with gr.Column(scale=5):
            gr.Markdown("# 🏆 VizWiz Benchmark Arena")
            gr.Markdown(
                "Automated evaluation platform for VizWiz challenges — "
                "datasets collected from blind photographers using a smartphone app."
            )
        with gr.Column(scale=1, min_width=160):
            gr.LoginButton(size="lg")
            home_greeting = gr.Markdown("👋 Log in to submit.")

    gr.Markdown("---")
    gr.Markdown("## Challenges")
    gr.Markdown(
        "Choose a challenge to view its leaderboard, submit predictions, and track your results."
    )
    gr.HTML(_challenge_cards_html())
    gr.Markdown(
        "---\n*More challenges coming soon. "
        "All challenges use HuggingFace OAuth — log in once to access everything.*"
    )

    demo.load(get_user_greeting, outputs=[home_greeting])

# ── OBJECT LOCALIZATION PAGE ──────────────────────────────────────────
with demo.route("Object Localization", "/object-localization") as obj_loc:

    with gr.Row():
        with gr.Column(scale=5):
            gr.Markdown("# 🎯 Object Localization Challenge")
            gr.Markdown(
                "Submit bounding box and instance segmentation predictions "
                "evaluated automatically against hidden ground-truth annotations."
            )
        with gr.Column(scale=1, min_width=160):
            gr.LoginButton(size="lg")
            ol_greeting = gr.Markdown("👋 Log in to submit.")

    gr.Markdown("---")

    with gr.Tabs():

        # Leaderboard
        with gr.TabItem("🏆 Leaderboard"):
            gr.Markdown("### Standard Phase Rankings")
            gr.Markdown(f"Ranked by **{DEFAULT_SORT_METRIC}** (descending). Standard phase only.")
            with gr.Accordion("📐 How is the Score Calculated?", open=False):
                gr.Markdown(EVAL_DETAILS_MD)
            lb_msg = gr.Markdown("")
            lb_table = gr.Dataframe(interactive=False, wrap=True)
            refresh_lb_btn = gr.Button("🔄 Refresh Leaderboard", variant="secondary", size="sm")

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

            refresh_lb_btn.click(refresh_leaderboard, outputs=[lb_table, lb_msg, ol_greeting])
            obj_loc.load(refresh_leaderboard, outputs=[lb_table, lb_msg, ol_greeting])

        # Submit
        with gr.TabItem("🚀 Submit Predictions"):
            with gr.Row():
                submit_greeting = gr.Markdown("👋 Log in with HuggingFace to submit.")
                cap_info = gr.Markdown("")
            gr.Markdown("---")
            with gr.Row():
                with gr.Column(scale=3):
                    gr.Markdown("#### Upload Submission File")
                    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")
                    team_input = gr.Textbox(label="Team / Display Name", placeholder="e.g. My Awesome Team")
                    model_input = gr.Textbox(label="Model Name", placeholder="e.g. ResNet50-FPN")
                    phase_input = gr.Dropdown(
                        label="Phase",
                        choices=[p["label"] for p in PHASES],
                        value=PHASES[0]["label"],
                    )
                    challenge_input = gr.Radio(
                        label="Challenge Type",
                        choices=CHALLENGE_TYPES,
                        value=CHALLENGE_TYPES[0],
                    )
            submit_btn = gr.Button("Submit (Queue for Evaluation)", variant="primary", size="lg")
            submit_status = gr.Markdown("")
            submission_id_box = gr.Code(label="Submission ID", language=None, visible=False)

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

            submit_btn.click(
                do_submit,
                inputs=[file_input, team_input, model_input, phase_input, challenge_input],
                outputs=[submit_status, submission_id_box],
            )

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

            obj_loc.load(update_submit_ui, outputs=[submit_greeting, cap_info])

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

            def refresh_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)

            refresh_my_btn.click(
                refresh_my_subs,
                inputs=[phase_filter],
                outputs=[my_sub_table, my_sub_msg, my_sub_stats, my_sub_greeting],
            )
            phase_filter.change(
                refresh_my_subs,
                inputs=[phase_filter],
                outputs=[my_sub_table, my_sub_msg, my_sub_stats, my_sub_greeting],
            )
            obj_loc.load(
                refresh_my_subs,
                inputs=[phase_filter],
                outputs=[my_sub_table, my_sub_msg, my_sub_stats, my_sub_greeting],
            )

# ── VQA PAGE ──────────────────────────────────────────────────────────
with demo.route("Visual Question Answering", "/vqa"):
    gr.Markdown("# 🤔 Visual Question Answering")
    gr.Markdown("### 🔒 Coming Soon")
    gr.Markdown(
        "This challenge is currently under development. "
        "Check back soon or follow [@VizWiz](https://vizwiz.org) for updates."
    )

# ── ANSWER GROUNDING PAGE ─────────────────────────────────────────────
with demo.route("Answer Grounding", "/answer-grounding"):
    gr.Markdown("# 📍 Answer Grounding")
    gr.Markdown("### 🔒 Coming Soon")
    gr.Markdown(
        "This challenge is currently under development. "
        "Check back soon or follow [@VizWiz](https://vizwiz.org) for updates."
    )


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