main.py

"""
PGSPD Main Entry Point
======================
Run the full pipeline for one or more paper subfolders under data/.

Usage: python main.py
Edit the PAPERS list below to specify which papers to run.
"""

from __future__ import annotations

import asyncio
import json
import logging
import sys
import traceback
from pathlib import Path

from pgspd import get_version
from pgspd.config import PGSPDConfig
from pgspd.pipeline import PGSPDPipeline, PipelineResult

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logging.getLogger("pyan.analyzer").setLevel(logging.WARNING)
logger = logging.getLogger("main")

DATA_DIR = Path(__file__).parent / "data"
PAPER_FILE = "paper.txt"
PAPER_FILE_FALLBACK = "paper.md"
CODEBASE_DIR = "codebase"
TASK_FILE = "task.json"  # expected keys: latex_code, task_id (optional)

# ↓↓↓ Specify which papers to run ↓↓↓
PAPERS = [
    "15-RAPTOR_RECURSIVE_ABSTRACTIVE_PROCESSING_FOR_TREE-ORGANIZED_RETRIEVAL",
]


def resolve_paper(name: str) -> tuple[Path, Path, str, str]:
    """Return (paper_txt, codebase, latex_code, task_id)."""
    paper_dir = DATA_DIR / name
    codebase = paper_dir / CODEBASE_DIR
    if not paper_dir.exists():
        raise FileNotFoundError(f"Paper folder not found: {paper_dir}")
    paper_txt = paper_dir / PAPER_FILE
    if not paper_txt.exists():
        paper_txt = paper_dir / PAPER_FILE_FALLBACK
    if not paper_txt.exists():
        raise FileNotFoundError(f"Missing {PAPER_FILE} (or {PAPER_FILE_FALLBACK}) in {paper_dir}")
    if not codebase.exists():
        raise FileNotFoundError(f"Missing {CODEBASE_DIR}/ in {paper_dir}")

    # Load latex_code from task.json if present
    task_json = paper_dir / TASK_FILE
    latex_code = ""
    task_id = name
    if task_json.exists():
        task_data = json.loads(task_json.read_text(encoding="utf-8"))
        latex_code = task_data.get("latex_code", "")
        task_id = task_data.get("task_id", name)
    else:
        logger.warning("No %s found in %s — latex_code will be empty (Step 1 will fail)", TASK_FILE, paper_dir)

    return paper_txt, codebase, latex_code, task_id


def make_config(paper_dir: Path) -> PGSPDConfig:
    cfg = PGSPDConfig()
    cfg.output_dir = str(paper_dir)
    cfg.save_intermediates = True
    return cfg


def save_result_summary(result: PipelineResult, paper_dir: Path) -> None:
    mods = result.spec.module_decomposition.modules if result.spec else []
    summary = {
        "paper": result.paper_path,
        "repo": result.repo_path,
        "steps_completed": result.step_completed,
        "task_id": mods[0].task_id if mods else "",
        "task_module": mods[0].name if mods else "",
        "code_nodes": result.code_graph_summary.get("num_nodes", 0),
        "code_edges": result.code_graph_summary.get("num_edges", 0),
        "modules_matched": sum(1 for m in result.matched_modules if m.entry_fqn),
        "modules_total": len(result.matched_modules),
        "closure_dirs": result.closure_dirs,
    }
    out = paper_dir / "pipeline_result.json"
    out.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    logger.info("Result summary → %s", out)


async def run_paper(name: str, paper_txt: Path, codebase: Path, latex_code: str, task_id: str) -> PipelineResult | None:
    banner = f" {name} "
    logger.info("=" * 60)
    logger.info(banner.center(60, "="))
    logger.info("=" * 60)

    pipeline = PGSPDPipeline(
        paper_path=str(paper_txt),
        repo_path=str(codebase),
        config=make_config(paper_txt.parent),
        latex_code=latex_code,
        task_id=task_id,
    )
    try:
        result = await pipeline.run(stop_after_step=4)
        save_result_summary(result, paper_txt.parent)
        logger.info("✓ '%s' done — %d/4 steps completed", name, result.step_completed)
        return result
    except Exception:
        logger.error("✗ '%s' failed:\n%s", name, traceback.format_exc())
        return None


async def main() -> None:
    logger.info("PGSPD version: %s", get_version())
    logger.info("Papers to run: %s", PAPERS)

    results: list[tuple[str, PipelineResult | None]] = []
    for name in PAPERS:
        try:
            paper_txt, codebase, latex_code, task_id = resolve_paper(name)
        except FileNotFoundError as e:
            logger.error("%s", e)
            results.append((name, None))
            continue
        results.append((name, await run_paper(name, paper_txt, codebase, latex_code, task_id)))

    logger.info("")
    logger.info("=" * 60)
    logger.info("ALL DONE")
    logger.info("=" * 60)
    for name, result in results:
        if result is None:
            logger.info("  %-36s  FAILED", name)
        else:
            n_dirs = len(result.closure_dirs)
            logger.info("  %-36s  %d/4 steps  closure_dirs=%d", name, result.step_completed, n_dirs)


if __name__ == "__main__":
    asyncio.run(main())