app.py

from __future__ import annotations

import json
from pathlib import Path

import gradio as gr


ROOT = Path(__file__).parent
MANIFEST_PATH = ROOT / "data" / "manifest.json"


CONDITION_LABELS = {
    "dataset": "Dataset trajectory",
    "draw": "Drawn trajectory",
    "specialist": "No trajectory input",
}


def load_manifest() -> dict:
    with MANIFEST_PATH.open(encoding="utf-8") as f:
        return json.load(f)


MANIFEST = load_manifest()
RUNS = MANIFEST["runs"]


def score_label(score: float | None) -> str:
    if score is None:
        return "n/a"
    return f"{score * 100:.1f}%"


def task_options() -> list[str]:
    return ["All"] + sorted({run["task"] for run in RUNS})


def model_options() -> list[str]:
    return ["All"] + sorted({run["model"] for run in RUNS})


def condition_options() -> list[str]:
    return ["All"] + sorted({run["condition"] for run in RUNS})


def matching_runs(task: str, model: str, condition: str) -> list[dict]:
    return [
        run
        for run in RUNS
        if (task == "All" or run["task"] == task)
        and (model == "All" or run["model"] == model)
        and (condition == "All" or run["condition"] == condition)
    ]


def run_label(run: dict) -> str:
    condition = CONDITION_LABELS.get(run["condition"], run["condition"])
    return f"{run['task']} | {run['model']} | {condition} | {score_label(run.get('mean_score'))}"


def summary_rows() -> list[list[str]]:
    rows = []
    for run in RUNS:
        rows.append(
            [
                run["task"],
                run["model"],
                CONDITION_LABELS.get(run["condition"], run["condition"]),
                score_label(run.get("mean_score")),
                str(run.get("num_episodes", "")),
                run.get("hf_eval_dataset_id") or "",
            ]
        )
    return rows


def update_run_choices(task: str, model: str, condition: str):
    runs = matching_runs(task, model, condition)
    choices = [run_label(run) for run in runs]
    value = choices[0] if choices else None
    max_episode = runs[0]["num_episodes"] if runs else 1
    return (
        gr.update(choices=choices, value=value),
        gr.update(minimum=1, maximum=max_episode, value=1),
    )


def selected_run(label: str | None) -> dict | None:
    if not label:
        return None
    for run in RUNS:
        if run_label(run) == label:
            return run
    return None


def render_episode(run_choice: str | None, episode_number: int):
    run = selected_run(run_choice)
    if run is None:
        return None, None, "No run selected.", "No trajectory input."

    episode_number = int(episode_number)
    episodes = run["episodes"]
    episode = next((ep for ep in episodes if ep["episode_number"] == episode_number), episodes[0])

    video = episode.get("video")
    overlay = episode.get("overlay")
    if overlay:
        overlay = str(ROOT / overlay)

    if run["condition"] == "specialist":
        trajectory_text = "Specialist baseline: no trajectory sketch or dataset trajectory was provided."
    elif episode.get("trajectory_source_mode") == "dataset":
        trajectory_text = f"Dataset trajectory from source episode {episode.get('trajectory_source_episode')}."
    elif episode.get("trajectory_source_mode") == "draw":
        trajectory_text = "Drawn trajectory input."
    else:
        trajectory_text = "Trajectory metadata unavailable."

    comment = episode.get("comment") or ""
    error = episode.get("episode_error") or ""
    details = "\n".join(
        line
        for line in [
            f"Task: {run['task']}",
            f"Model: {run['model']}",
            f"Condition: {CONDITION_LABELS.get(run['condition'], run['condition'])}",
            f"Run: {run.get('run_id')}",
            f"Episode: {episode.get('episode_number')}",
            f"Score: {score_label(episode.get('score'))}",
            f"Dataset: {run.get('hf_eval_dataset_id')}",
            f"Comment: {comment}" if comment else "",
            f"Error: {error}" if error else "",
            episode.get("video_note") or "The video is the full LeRobot chunk for this run. Use the slider to change the overlay image to the current dataset episode.",
            "Video unavailable in manifest." if not video else "",
        ]
        if line
    )
    return video, overlay, details, trajectory_text


def update_episode_bounds(run_choice: str | None):
    run = selected_run(run_choice)
    max_episode = run["num_episodes"] if run else 1
    return gr.update(minimum=1, maximum=max_episode, value=1)


css = """
.main-header h1 { margin-bottom: 0.25rem; }
.main-header p { margin-top: 0; color: #475569; }
textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace; }
"""


with gr.Blocks(title=MANIFEST.get("title", "Trajectory Showcase")) as demo:
    gr.Markdown(
        f"""
        <div class="main-header">
        <h1>{MANIFEST.get("title", "Diffusion Trajectory Model Showcase")}</h1>
        <p>Episode-level videos, trajectory inputs, and scores generated from experiment logs.</p>
        </div>
        """
    )

    gr.Dataframe(
        headers=["Task", "Model", "Condition", "Mean score", "Episodes", "Dataset"],
        value=summary_rows(),
        interactive=False,
        wrap=True,
    )

    with gr.Row():
        task_filter = gr.Dropdown(task_options(), value="All", label="Task")
        model_filter = gr.Dropdown(model_options(), value="All", label="Model")
        condition_filter = gr.Dropdown(condition_options(), value="All", label="Condition")

    initial_runs = matching_runs("All", "All", "All")
    run_select = gr.Dropdown(
        [run_label(run) for run in initial_runs],
        value=run_label(initial_runs[0]) if initial_runs else None,
        label="Run",
    )
    episode_slider = gr.Slider(
        minimum=1,
        maximum=initial_runs[0]["num_episodes"] if initial_runs else 1,
        value=1,
        step=1,
        label="Episode",
    )

    gr.Markdown("Use the slider to change the overlay image to the current dataset episode.")

    with gr.Row():
        video = gr.Video(label="Result video")
        overlay = gr.Image(label="Trajectory input", type="filepath")

    trajectory_text = gr.Textbox(label="Trajectory source", interactive=False)
    details = gr.Textbox(label="Episode metadata", lines=10, interactive=False)

    for control in (task_filter, model_filter, condition_filter):
        control.change(
            update_run_choices,
            inputs=[task_filter, model_filter, condition_filter],
            outputs=[run_select, episode_slider],
        ).then(render_episode, inputs=[run_select, episode_slider], outputs=[video, overlay, details, trajectory_text])

    run_select.change(update_episode_bounds, inputs=run_select, outputs=episode_slider).then(
        render_episode,
        inputs=[run_select, episode_slider],
        outputs=[video, overlay, details, trajectory_text],
    )
    episode_slider.change(
        render_episode,
        inputs=[run_select, episode_slider],
        outputs=[video, overlay, details, trajectory_text],
    )

    demo.load(render_episode, inputs=[run_select, episode_slider], outputs=[video, overlay, details, trajectory_text])


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