utils.py

"""
Evaluation utilities for the DFL Benchmark.

Storage: a private HF dataset repo holds tasks/<task>/Y_test.json and the
leaderboard results.json. Submitted JSON files are evaluated in-memory.

Env vars (HF Space secrets in production):
  HF_TOKEN         — read+write token for the private dataset
  HF_DATASET_REPO  — defaults to "GT-KOALA/DFL-Bench-Data"
"""

import importlib.util
import io
import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd
import requests
from huggingface_hub import HfApi, hf_hub_download, hf_hub_url

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "GT-KOALA/DFL-Bench-Data")
HF_TOKEN = os.environ.get("HF_TOKEN")
RESULTS_FILENAME = "results.json"

TASKS = [
    {"key": "power_scheduling", "label": "Power Scheduling", "icon": "⚡"},
    {"key": "newsvendor",       "label": "Newsvendor",       "icon": "📦"},
]
TASK_KEYS = [t["key"] for t in TASKS]
TASK_BY_KEY = {t["key"]: t for t in TASKS}

_api = HfApi(token=HF_TOKEN)


def _y_test_remote_path(task_key: str) -> str:
    return f"tasks/{task_key}/Y_test.json"


def _load_eval_module(task_key: str):
    path = Path(__file__).parent / "tasks" / task_key / "eval.py"
    spec = importlib.util.spec_from_file_location(f"_eval_{task_key}", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


_EVAL_MODULES = {t["key"]: _load_eval_module(t["key"]) for t in TASKS}


# ---------------------------------------------------------------------------
# HF dataset I/O
# ---------------------------------------------------------------------------

def _fetch_text_fresh(filename: str) -> str | None:
    """
    Fetch a file's current contents from the dataset repo, bypassing both the
    local HF cache and the CDN edge cache.

    hf_hub_download(force_download=True) only skips the *local* cache; the
    `main` resolve URL is still served stale by the CDN for a short window
    after an upload. We hit the resolve URL directly with a unique query
    param + no-cache headers so a refresh always reflects the latest commit.
    """
    url = hf_hub_url(repo_id=HF_DATASET_REPO, filename=filename, repo_type="dataset")
    headers = {"Cache-Control": "no-cache", "Pragma": "no-cache"}
    if HF_TOKEN:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"
    try:
        resp = requests.get(
            url,
            headers=headers,
            params={"cb": int(time.time() * 1000)},  # cache-buster
            timeout=20,
        )
        if resp.status_code == 404:
            return None
        resp.raise_for_status()
        return resp.text
    except requests.RequestException as e:
        print(f"[hf-dataset] failed to fetch {filename} from {HF_DATASET_REPO}: {e}")
        return None


def _upload_bytes(filename: str, data: bytes, commit_message: str):
    """Upload bytes to a path in the private dataset repo."""
    _api.upload_file(
        path_or_fileobj=io.BytesIO(data),
        path_in_repo=filename,
        repo_id=HF_DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN,
        commit_message=commit_message,
    )


# ---------------------------------------------------------------------------
# Results persistence
# ---------------------------------------------------------------------------

def load_results() -> list[dict]:
    """Load all evaluation results from the HF dataset repo (always fresh)."""
    text = _fetch_text_fresh(RESULTS_FILENAME)
    if text is None:
        return []
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return []


def save_result(entry: dict):
    """Append a result entry and upload the updated results.json."""
    results = load_results()
    results.append(entry)
    payload = json.dumps(results, indent=2).encode("utf-8")
    _upload_bytes(
        RESULTS_FILENAME,
        payload,
        commit_message=(
            f"Add submission: {entry.get('team', '?')} / "
            f"{entry.get('model', '?')} / {entry.get('task', '?')}"
        ),
    )


_LB_COLUMNS = ["Rank", "Team", "Model", "Regret ↓", "Constraint Viol. ↓"]


def load_results_as_dataframe(task_filter: str | None = None) -> pd.DataFrame:
    """Return results as a sorted DataFrame for the leaderboard, optionally filtered by task."""
    results = load_results()
    if task_filter:
        results = [r for r in results if r.get("task") == task_filter]
    if not results:
        return pd.DataFrame(columns=_LB_COLUMNS)

    df = pd.DataFrame(results)
    # Task loss is the decision-focused metric — sort ascending by its mean.
    df = df.sort_values("task_loss_mean", ascending=True).reset_index(drop=True)
    df.index += 1

    def _mean_std(mean_col: str, std_col: str) -> list[str]:
        if std_col in df.columns:
            return [
                f"{m:.4f} ± {s:.4f}" if pd.notna(s) else f"{m:.4f}"
                for m, s in zip(df[mean_col], df[std_col])
            ]
        return [f"{m:.4f}" for m in df[mean_col]]

    return pd.DataFrame({
        "Rank":               list(df.index),
        "Team":               df["team"].tolist(),
        "Model":              df["model"].tolist(),
        "Regret ↓":           _mean_std("task_loss_mean", "task_loss_std"),
        "Constraint Viol. ↓": _mean_std("constraint_violation_mean", "constraint_violation_std"),
    })


def load_task_figure(task_filter: str | None = None):
    """Scatter of Task Loss mean (y) vs std (x); constraint violators in red."""
    import plotly.graph_objects as go

    results = load_results()
    if task_filter:
        results = [r for r in results if r.get("task") == task_filter]

    fig = go.Figure()
    feasible = [r for r in results if (r.get("constraint_violation_mean") or 0) <= 0]
    violators = [r for r in results if (r.get("constraint_violation_mean") or 0) > 0]

    for group, color, name in (
        (feasible, "#2563eb", "Feasible"),
        (violators, "#dc2626", "Constraint violation"),
    ):
        if not group:
            continue
        fig.add_trace(go.Scatter(
            x=[r.get("task_loss_std", 0.0) for r in group],
            y=[r.get("task_loss_mean", 0.0) for r in group],
            mode="markers",
            name=name,
            marker=dict(color=color, size=11, line=dict(width=0.5, color="#ffffff")),
            customdata=[[r.get("model", "?"), r.get("team", "?")] for r in group],
            hovertemplate=(
                "Method: %{customdata[0]}<br>"
                "Group: %{customdata[1]}<br>"
                "Regret: %{y:.4f}<br>"
                "Std: %{x:.4f}<extra></extra>"
            ),
        ))

    fig.update_layout(
        xaxis_title="Regret Std",
        yaxis_title="Regret Mean",
        margin=dict(l=50, r=20, t=20, b=45),
        height=340,
        legend=dict(orientation="h", yanchor="bottom", y=1.02, x=0),
        plot_bgcolor="#ffffff",
        paper_bgcolor="#ffffff",
        font=dict(size=12),
    )
    fig.update_xaxes(showgrid=True, gridcolor="#f3f4f6", zeroline=False)
    fig.update_yaxes(showgrid=True, gridcolor="#f3f4f6", zeroline=False)
    return fig


def latest_submission_time(task_filter: str | None = None) -> str | None:
    """ISO timestamp of the most recent submission (optionally for one task), or None."""
    results = load_results()
    if task_filter:
        results = [r for r in results if r.get("task") == task_filter]
    times = [r.get("submitted_at") for r in results if r.get("submitted_at")]
    return max(times) if times else None


def format_time_ago(iso_str: str | None) -> str:
    """Humanize an ISO timestamp as 'N hours ago' / 'N days ago' etc."""
    if not iso_str:
        return "no submissions yet"
    try:
        dt = datetime.fromisoformat(iso_str.replace("Z", "+00:00"))
    except (ValueError, TypeError):
        return "unknown"
    delta = datetime.now(timezone.utc) - dt
    seconds = int(delta.total_seconds())
    if seconds < 60:
        return "just now"
    if seconds < 3600:
        m = seconds // 60
        return f"{m} minute{'s' if m != 1 else ''} ago"
    if seconds < 86400:
        h = seconds // 3600
        return f"{h} hour{'s' if h != 1 else ''} ago"
    d = seconds // 86400
    return f"{d} day{'s' if d != 1 else ''} ago"


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def evaluate_predictions(submission_path: str, task_key: str) -> dict:
    """Run the task's evaluator against Y_test.json from the data repo."""
    if task_key not in _EVAL_MODULES:
        raise ValueError(f"Unknown task: {task_key}")
    y_test_path = hf_hub_download(
        repo_id=HF_DATASET_REPO,
        filename=_y_test_remote_path(task_key),
        repo_type="dataset",
        token=HF_TOKEN,
    )
    return _EVAL_MODULES[task_key].evaluate(submission_path, y_test_path)


# ---------------------------------------------------------------------------
# Full submission handler
# ---------------------------------------------------------------------------

def handle_submission(team_name: str, model_name: str, task_key: str, file_obj) -> str:
    """Validate → evaluate JSON → persist metrics → return status."""
    if not team_name or not team_name.strip():
        return "❌ Please enter a team name."
    if not model_name or not model_name.strip():
        return "❌ Please enter a model name."
    if task_key not in TASK_KEYS:
        return f"❌ Unknown task: {task_key}"
    if file_obj is None:
        return "❌ Please upload a submission JSON file."

    try:
        scores = evaluate_predictions(file_obj, task_key)
    except Exception as e:
        return f"❌ Evaluation error: {e}"

    entry = {
        "team": team_name.strip(),
        "model": model_name.strip(),
        "task": task_key,
        "task_loss_mean": scores["task_loss_mean"],
        "task_loss_std": scores["task_loss_std"],
        "constraint_violation_mean": scores["constraint_violation_mean"],
        "constraint_violation_std": scores["constraint_violation_std"],
        "submitted_at": datetime.now(timezone.utc).isoformat(),
    }
    try:
        save_result(entry)
    except Exception as e:
        return f"❌ Could not persist result: {e}"

    return (
        "✅ **Submission evaluated successfully!**\n\n"
        "| Metric | Mean ± Std |\n"
        "|--------|------------|\n"
        f"| **Regret** | **{scores['task_loss_mean']:.4f} ± {scores['task_loss_std']:.4f}** |\n"
        f"| Constraint Violation | {scores['constraint_violation_mean']:.4f} ± {scores['constraint_violation_std']:.4f} |\n"
    )