scripts/plan.py

"""Compute a prerequisite reading path for a target RL paper.

Usage:
    uv run python scripts/plan.py --target 1710.02298                   # plan for Rainbow, no prior knowledge
    uv run python scripts/plan.py --target 1710.02298 --known 1312.5602  # already know DQN
    uv run python scripts/plan.py --target 1710.02298 --known DQN,A3C   # use tags or arxiv IDs
    uv run python scripts/plan.py --list                                # list all canonical papers
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from rich.console import Console
from rich.panel import Panel

from researchpath.corpus import CANONICAL_RL_PAPERS
from researchpath.planning import plan_reading_path

console = Console()
_TAG_TO_ID = {p.tag.lower(): p.arxiv_id for p in CANONICAL_RL_PAPERS}
_ID_TO_PAPER = {p.arxiv_id: p for p in CANONICAL_RL_PAPERS}


def _resolve(token: str) -> str:
    """Accept either an arxiv ID or a paper tag (case-insensitive)."""
    token = token.strip()
    if token in _ID_TO_PAPER:
        return token
    if token.lower() in _TAG_TO_ID:
        return _TAG_TO_ID[token.lower()]
    raise SystemExit(f"Unknown paper: {token!r}. Use --list to see options.")


def _list_papers() -> None:
    console.print("[bold cyan]Canonical RL papers:[/bold cyan]\n")
    for p in CANONICAL_RL_PAPERS:
        console.print(f"  [{p.arxiv_id}]  [bold]{p.tag:20s}[/bold]  ({p.year})  {p.title}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Plan a prerequisite reading path.")
    parser.add_argument("--target", help="arxiv ID or tag of the paper you want to read.")
    parser.add_argument(
        "--known",
        default="",
        help="Comma-separated arxiv IDs or tags you already understand (default: none).",
    )
    parser.add_argument("--list", action="store_true", help="List all canonical papers.")
    args = parser.parse_args()

    if args.list:
        _list_papers()
        return 0

    if not args.target:
        parser.error("--target is required (or pass --list).")

    target_id = _resolve(args.target)
    known_ids = {_resolve(t) for t in args.known.split(",") if t.strip()}

    plan = plan_reading_path(target_id, known_ids=known_ids)
    target = _ID_TO_PAPER[target_id]

    header = f"Reading path for [{target.arxiv_id}] {target.tag} — {target.title}"
    if known_ids:
        known_tags = ", ".join(sorted(_ID_TO_PAPER[k].tag for k in known_ids))
        header += f"\nAssuming you already know: {known_tags}"

    console.print(Panel(header, border_style="cyan"))
    console.print()
    console.print(plan.render_text())
    console.print()
    console.print(
        f"[dim]{len(plan.steps)} step(s). Static prerequisite graph; "
        f"add LLM-extracted concept dependencies in a future iteration.[/dim]"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())