app.py

"""Gradio Blocks dashboard for AgentAlign Lab.

Tabs:
1. Overview — baseline vs tuned metrics, bar chart of task success rates
2. Task Explorer — browse tasks, view files and verifier config
3. Trajectory Viewer — step-by-step trajectory with verifier score
4. Preference Pair Viewer — side-by-side chosen vs rejected with score margin
5. Baseline vs Tuned — model comparison table
6. Failure Analysis — failure label distribution and example runs

Data is loaded from data/ and outputs/ directories on startup.
"""

import json
from collections import Counter, defaultdict
from pathlib import Path

from agentalign.data.trajectories import load_trajectories
from agentalign.schemas import PreferencePair
from agentalign.tasks.load import load_task

# ---------------------------------------------------------------------------
# Data directories
# ---------------------------------------------------------------------------

TASKS_DIR = Path("data/tasks")
PRIMARY_SCORED_DIRS = [
    Path("data/trajectories/scored_train"),
    Path("data/trajectories/scored_val"),
    Path("data/trajectories/scored_test"),
    Path("data/trajectories/scored_llm_train"),
    Path("data/trajectories/scored_llm_val"),
    Path("data/trajectories/scored_llm_test"),
]
LEGACY_SCORED_DIR = Path("data/trajectories/scored")
PREF_PATHS = [
    Path("data/preferences/dpo_train.jsonl"),
    Path("data/preferences/dpo_val.jsonl"),
    Path("data/preferences/dpo_llm_train.jsonl"),
    Path("data/preferences/dpo_llm_val.jsonl"),
]


# ---------------------------------------------------------------------------
# Data loaders
# ---------------------------------------------------------------------------

def _load_tasks():
    """Load all individual task JSON files."""
    if not TASKS_DIR.exists():
        return []
    return [load_task(p) for p in sorted(TASKS_DIR.glob("*.json"))]


def _load_trajectories():
    """Load all scored trajectories from available directories."""
    trajectories = []
    scored_dirs = [p for p in PRIMARY_SCORED_DIRS if p.exists()]
    if not scored_dirs and LEGACY_SCORED_DIR.exists():
        scored_dirs = [LEGACY_SCORED_DIR]
    for scored_dir in scored_dirs:
        for path in sorted(scored_dir.glob("*.jsonl")):
            trajectories.extend(load_trajectories(path))
    return trajectories


def _load_preferences():
    """Load all preference pairs."""
    pairs = []
    for pref_path in PREF_PATHS:
        if not pref_path.exists():
            continue
        for line in pref_path.read_text().splitlines():
            if line.strip():
                pairs.append(PreferencePair.model_validate_json(line))
    return pairs


# ---------------------------------------------------------------------------
# Tab content functions
# ---------------------------------------------------------------------------

def _summarize_overview():
    """Build overview summary text."""
    tasks = _load_tasks()
    trajectories = _load_trajectories()
    preferences = _load_preferences()
    passed = sum(1 for t in trajectories if t.verifier_result and t.verifier_result.passed)
    family_counts = Counter(t.family for t in tasks)
    pass_rate = (passed / len(trajectories) * 100) if trajectories else 0.0
    lines = [
        f"Total tasks: {len(tasks)}",
        f"Task families: {dict(family_counts)}",
        f"Total trajectories: {len(trajectories)}",
        f"Total preference pairs: {len(preferences)}",
        f"Pass rate: {pass_rate:.1f}%",
    ]
    return "\n".join(lines)


def _task_table():
    return [
        [t.task_id, t.family, t.difficulty or "", t.verifier.type, len(t.files)]
        for t in _load_tasks()
    ]


def _task_ids():
    return [t.task_id for t in _load_tasks()]


def _trajectory_ids():
    return [t.run_id for t in _load_trajectories()]


def _preference_ids():
    return [p.pair_id for p in _load_preferences()]


def _show_task(task_id):
    tasks = {t.task_id: t for t in _load_tasks()}
    task = tasks.get(task_id)
    if not task:
        return "", "", {}
    files = "\n\n".join(
        f"### {f.path}\n```\n{f.content}\n```" for f in task.files
    )
    verifier = task.verifier.model_dump(mode="json")
    return task.instruction, files, verifier


def _show_trajectory(run_id):
    trajectories = {t.run_id: t for t in _load_trajectories()}
    traj = trajectories.get(run_id)
    if not traj:
        return "", {}, ""
    rows = [
        [
            step.step_index, step.thought, step.action,
            json.dumps(step.args, sort_keys=True),
            step.observation or "", step.error or "",
        ]
        for step in traj.steps
    ]
    verifier = traj.verifier_result.model_dump(mode="json") if traj.verifier_result else {}
    return rows, verifier, traj.final_answer or ""


def _show_preference(pair_id):
    pairs = {p.pair_id: p for p in _load_preferences()}
    pair = pairs.get(pair_id)
    if not pair:
        return "", "", "", ""
    scores = (
        f"chosen_score={pair.chosen_score}\n"
        f"rejected_score={pair.rejected_score}\n"
        f"margin={pair.score_margin}"
    )
    return pair.prompt, pair.chosen, pair.rejected, scores


def _model_comparison_table():
    by_agent = defaultdict(list)
    for traj in _load_trajectories():
        by_agent[traj.agent_id or traj.model or "unknown"].append(traj)
    rows = []
    for agent, trajectories in sorted(by_agent.items()):
        total = len(trajectories)
        passed = sum(1 for t in trajectories if t.verifier_result and t.verifier_result.passed)
        avg_score = (
            sum(t.verifier_result.score for t in trajectories if t.verifier_result) / total
            if total else 0.0
        )
        avg_steps = sum(len(t.steps) for t in trajectories) / total if total else 0.0
        rows.append([
            agent, total,
            f"{passed / total * 100:.1f}%",
            f"{avg_score:.2f}",
            f"{avg_steps:.1f}",
        ])
    return rows


def _failure_table():
    counter = Counter()
    examples = {}
    for traj in _load_trajectories():
        if not traj.verifier_result or traj.verifier_result.passed:
            continue
        tags = traj.verifier_result.failure_tags or ["unknown_failure"]
        for tag in tags:
            counter[tag] += 1
            examples.setdefault(tag, traj.run_id)
    return [[tag, count, examples[tag]] for tag, count in counter.most_common()]


def _side_by_side_task_ids():
    """Return task_ids that have trajectories from multiple agents."""
    by_task: dict[str, set[str]] = defaultdict(set)
    for traj in _load_trajectories():
        agent = traj.agent_id or traj.model or "unknown"
        by_task[traj.task_id].add(agent)
    return sorted(tid for tid, agents in by_task.items() if len(agents) > 1)


def _show_side_by_side(task_id):
    """Show trajectories from different agents on the same task."""
    by_agent: dict[str, list] = defaultdict(list)
    for traj in _load_trajectories():
        if traj.task_id == task_id:
            agent = traj.agent_id or traj.model or "unknown"
            by_agent[agent].append(traj)

    lines: list[str] = []
    for agent in sorted(by_agent):
        trajs = by_agent[agent]
        t = trajs[0]
        passed = t.verifier_result.passed if t.verifier_result else False
        score = t.verifier_result.score if t.verifier_result else 0
        unsafe = t.verifier_result.unsafe_actions if t.verifier_result else 0
        n_steps = len(t.steps)
        status = "PASS" if passed else "FAIL"

        lines.append(
            f"### {agent} — {status} "
            f"(score={score:.1f}, steps={n_steps}, unsafe={unsafe})"
        )
        for step in t.steps:
            action_str = f"`{step.action}`"
            args_short = json.dumps(step.args, sort_keys=True)[:80] if step.args else ""
            thought_short = (step.thought or "")[:60]
            lines.append(
                f"- Step {step.step_index}: {thought_short}… → "
                f"{action_str}({args_short})"
            )
            obs = (step.observation or "")[:120]
            if obs:
                lines.append(f"  > {obs}")
        lines.append("")

    return "\n".join(lines) if lines else "No trajectories for this task."


def _eval_results_summary():
    """Load eval_results.json and format as a summary string."""
    eval_path = Path("outputs/evals/eval_results.json")
    if not eval_path.exists():
        return "No evaluation results found. Run the GPU eval notebook first."

    results = json.loads(eval_path.read_text())
    lines = ["## Evaluation Results\n"]

    for agent in ["qwen_base", "qwen_sft", "qwen_dpo"]:
        if agent not in results:
            continue
        m = results[agent]
        lines.append(
            f"**{agent}**: {m['passed']}/{m['total']} pass "
            f"({m['pass_rate']*100:.1f}%), "
            f"unsafe={m.get('unsafe_actions', 0)}, "
            f"avg_steps={m.get('avg_steps', 0)}"
        )

    lines.append("\n### Bootstrap 95% CIs")
    for key, label in [
        ("dpo_vs_base", "DPO − Base"),
        ("sft_vs_base", "SFT − Base"),
        ("dpo_vs_sft", "DPO − SFT"),
    ]:
        if key in results:
            r = results[key]
            spans = (
                "spans zero"
                if r["ci_lo"] <= 0 <= r["ci_hi"]
                else "significant"
            )
            lines.append(
                f"- **{label}**: {r['diff']:+.4f} "
                f"[{r['ci_lo']:+.4f}, {r['ci_hi']:+.4f}] "
                f"p={r.get('p_value', 'N/A')} ({spans})"
            )

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# App builder
# ---------------------------------------------------------------------------

def build_app():
    """Build and return the Gradio Blocks dashboard."""
    import gradio as gr

    with gr.Blocks(title="AgentAlign Dashboard") as app:
        gr.Markdown("# AgentAlign Lab Dashboard")

        # Tab 1: Overview
        with gr.Tab("Overview"):
            refresh_btn = gr.Button("Refresh")
            summary = gr.Textbox(label="Metrics Summary", lines=8)
            refresh_btn.click(_summarize_overview, outputs=summary)
            app.load(_summarize_overview, outputs=summary)

        # Tab 2: Task Explorer
        with gr.Tab("Task Explorer"):
            gr.Dataframe(
                headers=["task_id", "family", "difficulty", "verifier", "files"],
                value=_task_table(),
                interactive=False,
            )
            task_selector = gr.Dropdown(label="Task", choices=_task_ids())
            instruction_box = gr.Textbox(label="Instruction", lines=3)
            files_md = gr.Markdown()
            verifier_json = gr.JSON(label="Verifier")
            task_selector.change(
                _show_task,
                inputs=task_selector,
                outputs=[instruction_box, files_md, verifier_json],
            )

        # Tab 3: Trajectory Viewer
        with gr.Tab("Trajectory Viewer"):
            traj_selector = gr.Dropdown(label="Trajectory", choices=_trajectory_ids())
            steps_df = gr.Dataframe(
                headers=["step", "thought", "action", "args", "observation", "error"],
                interactive=False,
            )
            vr_json = gr.JSON(label="Verifier Result")
            final_ans = gr.Textbox(label="Final Answer", lines=2)
            traj_selector.change(
                _show_trajectory,
                inputs=traj_selector,
                outputs=[steps_df, vr_json, final_ans],
            )

        # Tab 4: Preference Pair Viewer
        with gr.Tab("Preference Pair Viewer"):
            pref_selector = gr.Dropdown(label="Preference Pair", choices=_preference_ids())
            prompt_box = gr.Textbox(label="Prompt", lines=8)
            chosen_box = gr.Textbox(label="Chosen", lines=8)
            rejected_box = gr.Textbox(label="Rejected", lines=8)
            scores_box = gr.Textbox(label="Scores", lines=3)
            pref_selector.change(
                _show_preference,
                inputs=pref_selector,
                outputs=[prompt_box, chosen_box, rejected_box, scores_box],
            )

        # Tab 5: Baseline vs Tuned (enhanced)
        with gr.Tab("Baseline vs Tuned"):
            eval_md = gr.Markdown()
            app.load(_eval_results_summary, outputs=eval_md)
            gr.Dataframe(
                headers=["agent", "runs", "pass_rate", "avg_score", "avg_steps"],
                value=_model_comparison_table(),
                interactive=False,
            )

            # Side-by-side trajectory comparison
            gr.Markdown("### Side-by-Side Trajectory Comparison")
            sbs_selector = gr.Dropdown(
                label="Task (with multiple agents)",
                choices=_side_by_side_task_ids(),
            )
            sbs_output = gr.Markdown()
            sbs_selector.change(
                _show_side_by_side,
                inputs=sbs_selector,
                outputs=sbs_output,
            )

        # Tab 6: Failure Analysis
        with gr.Tab("Failure Analysis"):
            gr.Dataframe(
                headers=["failure_tag", "count", "example_run"],
                value=_failure_table(),
                interactive=False,
            )

    return app


try:
    import spaces

    @spaces.GPU
    def _zerogpu_noop():
        return None
except ImportError:
    pass


if __name__ == "__main__":
    app = build_app()
    app.launch()