"""Public leaderboard + submission UI for the OOD Reasoning Challenge.

Two tracks, two leaderboards:
  - track1 (Chain-of-Causation Generation): 6 rollouts/sample, top1/avgk/topk.
  - track2 (Reasoning Auto-Labeling):        1 rollout/sample, top1 only.

On submit, this space calls the PRIVATE ZeroGPU evaluator
(`nvidia/PhysicalAI-OOD-Evaluator-Private`) via gradio_client, using a token
that can access that private space. Scores come back in the same session; if the
user opted in, the evaluator also appends the row to the private results dataset,
which this space reads to render the leaderboard tables.

Secret (set in Space settings):
  HF_TOKEN   token that can CALL the private backend space and READ the results
             dataset (`nvidia/PhysicalAI-OOD-Results`).
"""

from __future__ import annotations

import json
import os

import gradio as gr
import pandas as pd
from gradio_client import Client, handle_file
from huggingface_hub import hf_hub_download

HF_TOKEN = os.environ["HF_TOKEN"]
BACKEND = os.environ.get("BACKEND_SPACE", "nvidia/PhysicalAI-OOD-Evaluator-Private")
RESULTS_REPO = os.environ.get("RESULTS_REPO", "nvidia/PhysicalAI-OOD-Results")
LEADERBOARD_FILE = "leaderboard.jsonl"
INFO_URL = "https://huggingface.co/spaces/nvidia/PhysicalAI-AV-OOD-Reasoning-Challenge-2026"

# Per-track display columns (both ranked by top1_score).
COLS = {
    "track1": ["Rank", "Model", "Organization", "top1_score", "avgk_score",
               "topk_score", "Coverage", "Submitted"],
    "track2": ["Rank", "Model", "Organization", "top1_score", "Coverage", "Submitted"],
}


def load_leaderboard(track: str) -> pd.DataFrame:
    """Build the ranked, published-only table for one track from the results dataset."""
    cols = COLS[track]
    try:
        path = hf_hub_download(
            RESULTS_REPO, LEADERBOARD_FILE, repo_type="dataset", token=HF_TOKEN
        )
        rows = [json.loads(ln) for ln in open(path).read().splitlines() if ln.strip()]
    except Exception:
        rows = []
    # Rows missing "track" are legacy track1 submissions.
    rows = [r for r in rows if r.get("published") and r.get("track", "track1") == track]
    if not rows:
        return pd.DataFrame(columns=cols)

    # Keep the best (highest top1_score) entry per model name.
    best: dict[str, dict] = {}
    for r in rows:
        m = r.get("model", "?")
        if m not in best or r.get("top1_score", 0) > best[m].get("top1_score", 0):
            best[m] = r
    df = pd.DataFrame(sorted(best.values(), key=lambda r: r.get("top1_score", 0), reverse=True))

    df.insert(0, "Rank", range(1, len(df) + 1))
    df["Coverage"] = (df["coverage"] * 100).round(1).astype(str) + "%"
    df["Submitted"] = df["submitted_at"].str.slice(0, 10)
    df = df.rename(columns={"model": "Model", "organization": "Organization"})
    for c in ("top1_score", "avgk_score", "topk_score"):
        if c in df.columns:
            df[c] = df[c].round(4)
    return df[cols]


def submit(pred_file, model_name, organization, contact, track, publish,
           progress=gr.Progress()):
    if not model_name or not model_name.strip():
        raise gr.Error("Please enter a model name.")
    if pred_file is None:
        raise gr.Error("Please upload a submission.json file.")
    progress(0.15, desc="Connecting to evaluator…")
    client = Client(BACKEND, token=HF_TOKEN)
    progress(0.4, desc="Running GPU evaluation (this can take a minute)…")
    row, md = client.predict(
        handle_file(pred_file), model_name, organization or "", contact or "",
        bool(publish), track, api_name="/evaluate",
    )
    progress(1.0, desc="Done")
    # Refresh both boards (a submission only changes one, but refreshing both is cheap).
    return md, load_leaderboard("track1"), load_leaderboard("track2")


ABOUT_MD = f"""
## Physical AI AV — OOD Reasoning Challenge

Challenge overview, **submission tutorials** (format + scoring for each track),
timeline, and prize are on the challenge page:

**[Challenge info & submission tutorials →]({INFO_URL})**

Quick reminder: untick **“Publish to public leaderboard”** on the Submit tab to
get your scores back privately without adding a row to the public board.
"""


with gr.Blocks(title="OOD Reasoning Challenge Leaderboard", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚗 Physical AI AV — OOD Reasoning Challenge Leaderboard")
    with gr.Tabs():
        with gr.Tab("🏆 Track 1 · Reasoning"):
            tbl1 = gr.Dataframe(value=lambda: load_leaderboard("track1"),
                                interactive=False, wrap=True)
            gr.Button("🔄 Refresh").click(lambda: load_leaderboard("track1"), None, tbl1)
        with gr.Tab("🏆 Track 2 · Auto-Labeling"):
            tbl2 = gr.Dataframe(value=lambda: load_leaderboard("track2"),
                                interactive=False, wrap=True)
            gr.Button("🔄 Refresh").click(lambda: load_leaderboard("track2"), None, tbl2)
        with gr.Tab("📤 Submit"):
            with gr.Row():
                with gr.Column():
                    f_track = gr.Radio(
                        choices=[("Track 1 · Reasoning (6 rollouts)", "track1"),
                                 ("Track 2 · Auto-Labeling (1 rollout)", "track2")],
                        value="track1", label="Track")
                    f_pred = gr.File(label="submission.json", file_types=[".json"])
                    f_name = gr.Textbox(label="Model name *", placeholder="e.g. MyModel-7B")
                    f_org = gr.Textbox(label="Organization")
                    f_contact = gr.Textbox(label="Contact email (optional)")
                    f_pub = gr.Checkbox(value=True, label="Publish to public leaderboard")
                    btn = gr.Button("Submit & Evaluate", variant="primary")
                with gr.Column():
                    out_md = gr.Markdown("Results will appear here after evaluation.")
            btn.click(submit, [f_pred, f_name, f_org, f_contact, f_track, f_pub],
                      [out_md, tbl1, tbl2])
        with gr.Tab("ℹ️ About"):
            gr.Markdown(ABOUT_MD)

if __name__ == "__main__":
    demo.queue(max_size=32).launch()