"""Gradio leaderboard for the UQ competition (Hugging Face Space).
- Every submission is ADDED as its own row (no overwrite; multiple per team allowed).
- Baselines (Baseline 1, random) always shown, never stored, never deduped away.
- Submissions persist to a PRIVATE dataset so they survive restarts and are editable there.
- Admin accordion: remove a row by sid / remove a team by name / reset all (password-gated).
See README.md for setup."""
import os, datetime, uuid
import numpy as np, pandas as pd, gradio as gr
import scoring
from huggingface_hub import HfApi, hf_hub_download

# ---- config: set these as Space variables/secrets ----
LABELS_DATASET = os.environ.get("LABELS_DATASET", "")               # private dataset holding test_labels.csv
SUBS_DATASET   = os.environ.get("SUBS_DATASET", LABELS_DATASET)     # where submissions.csv is stored (can be the same)
HF_TOKEN       = os.environ.get("HF_TOKEN", "")                     # token with WRITE access to SUBS_DATASET
ADMIN_PW       = os.environ.get("ADMIN_PW", "")                     # password for the admin controls
LABELS_FILE, SUBS_FILE = "test_labels.csv", "submissions.csv"
BOARD_COLS = ["sid", "team", "submitted_at", "Brier", "AUROC", "Accuracy", "n"]

api = HfApi(token=HF_TOKEN) if HF_TOKEN else None

# ---- hidden labels ----
def _read_labels():
    if LABELS_DATASET and HF_TOKEN:
        p = hf_hub_download(LABELS_DATASET, LABELS_FILE, repo_type="dataset", token=HF_TOKEN)
    else:
        p = LABELS_FILE                                            # local fallback (private Space)
    return pd.read_csv(p).set_index("id")["correct"].astype(int)

LABELS = _read_labels()
BASE_RATE = float(LABELS.mean())

# ---- persistence: read/write submissions.csv in the private dataset (else local file) ----
def load_board():
    try:
        if SUBS_DATASET and HF_TOKEN:
            p = hf_hub_download(SUBS_DATASET, SUBS_FILE, repo_type="dataset",
                                token=HF_TOKEN, force_download=True)
        else:
            p = SUBS_FILE
        df = pd.read_csv(p)
    except Exception:
        df = pd.DataFrame(columns=BOARD_COLS)
    for c in BOARD_COLS:
        if c not in df.columns:
            df[c] = pd.Series(dtype=object)
    return df[BOARD_COLS]

def save_board(df):
    if SUBS_DATASET and HF_TOKEN:
        api.upload_file(path_or_fileobj=df.to_csv(index=False).encode(),
                        path_in_repo=SUBS_FILE, repo_id=SUBS_DATASET,
                        repo_type="dataset", commit_message="update submissions")
    else:
        df.to_csv(SUBS_FILE, index=False)

# ---- baselines: recomputed every render, never stored, never deduped ----
def baseline_rows():
    rng = np.random.default_rng(0); ids = LABELS.index
    out = []
    for name, p in {"🤖 Baseline 1": np.full(len(ids), BASE_RATE),
                    "🤖 random":      rng.random(len(ids))}.items():
        m = scoring.grade(pd.DataFrame({"id": ids.astype(str), "p_correct": p}), LABELS)
        out.append({"sid": "—", "team": name, "submitted_at": "baseline", **m})
    return pd.DataFrame(out)

def render():
    board = pd.concat([load_board(), baseline_rows()], ignore_index=True)
    board = board.sort_values("Brier", ascending=True).reset_index(drop=True)
    board.insert(0, "rank", np.arange(1, len(board) + 1))
    return board[["rank", "team", "Brier", "AUROC", "Accuracy", "submitted_at", "sid"]]

# ---- actions ----
def submit(team, file):
    if not team or not team.strip():
        return "⚠️ Enter a team name.", render()
    team = team.strip()
    if team.startswith("🤖"):
        return "⚠️ That prefix is reserved for baselines.", render()
    if file is None:
        return "⚠️ Upload your submission.csv.", render()
    try:
        m = scoring.grade(pd.read_csv(file.name), LABELS)
    except Exception as e:
        return f"❌ {e}", render()
    sid = uuid.uuid4().hex[:6]
    row = {"sid": sid, "team": team,
           "submitted_at": datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M"), **m}
    save_board(pd.concat([load_board(), pd.DataFrame([row])], ignore_index=True))
    return (f"✅ **{team}** — Brier **{m['Brier']}** (AUROC {m['AUROC']}, Acc {m['Accuracy']}). "
            f"Entry id `{sid}`. Submit as many times as you like.", render())

def admin_remove(target, pw):
    if not ADMIN_PW or pw != ADMIN_PW:
        return "⚠️ Wrong or unset admin password.", render()
    t = target.strip()
    board = load_board()
    keep = board[(board["sid"] != t) & (board["team"].astype(str).str.strip() != t)]
    save_board(keep)
    return f"Removed {len(board) - len(keep)} row(s) matching `{t}`.", render()

def admin_reset(pw):
    if not ADMIN_PW or pw != ADMIN_PW:
        return "⚠️ Wrong or unset admin password.", render()
    save_board(pd.DataFrame(columns=BOARD_COLS))
    return "Leaderboard cleared (baselines remain).", render()

with gr.Blocks(title="UQ Competition") as demo:
    gr.Markdown("# 🏆 UQ Competition Leaderboard\n"
                "Predict P(model's answer is correct). **Ranked by Brier — lower is better.** "
                "You may submit as many times as you like; **every submission is added** as its own row.")
    with gr.Row():
        team = gr.Textbox(label="Team name", scale=2)
        file = gr.File(label="submission.csv", file_types=[".csv"], scale=2)
    btn = gr.Button("Submit", variant="primary")
    status = gr.Markdown()
    table = gr.Dataframe(value=render(), interactive=False, label="Leaderboard")
    btn.click(submit, [team, file], [status, table])

    with gr.Accordion("admin", open=False):
        gr.Markdown("Remove one entry by its **sid** (last column), or remove all of a team by **name**.")
        tgt = gr.Textbox(label="sid or team name to remove")
        pw  = gr.Textbox(label="admin password", type="password")
        with gr.Row():
            gr.Button("Remove").click(admin_remove, [tgt, pw], [status, table])
            gr.Button("Reset ALL", variant="stop").click(admin_reset, [pw], [status, table])

    demo.load(render, None, table)

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