setup_papers.py

"""
Download abstracts and PDFs to reconstruct the paper store for EvidenceSynthBench.

Usage:
    python setup_papers.py download                    # Download abstracts + OA PDFs
    python setup_papers.py download --abstracts-only   # Abstracts only, skip PDF download
    python setup_papers.py parse                       # Parse downloaded PDFs with GROBID
    python setup_papers.py parse --grobid-url URL      # Custom GROBID URL
"""

import argparse
import json
import re
import sys
import time
from pathlib import Path

import requests

from synthesisbench.resolve import (
    API_DELAY,
    batch_fetch_abstracts_and_titles,
    check_openalex_oa,
    check_pmc_availability,
)
from synthesisbench.fulltext import check_grobid_available
from synthesisbench.grobid_parser import GrobidSectionParser, GrobidParsingError

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
PAPERS_DIR = DATA_DIR / "paper_store"
PDFS_DIR = DATA_DIR / "pdfs"
SUBSETS_DIR = DATA_DIR / "subsets"
STUDIES_RIS = DATA_DIR / "studies.ris"

SUBSET_CONFIG = {
    "rob": {
        "manifest": "rob.json",
        "task_dir": DATA_DIR / "tasks" / "rob",
        "fulltext_required": True,
    },
    "eligibility": {
        "manifest": "eligibility.json",
        "task_dir": DATA_DIR / "tasks" / "eligibility",
        "fulltext_required": True,
    },
    "directness": {
        "manifest": "directness.json",
        "task_dir": DATA_DIR / "tasks" / "directness",
        "fulltext_required": False,
    },
}


# ---------------------------------------------------------------------------
# PMID discovery
# ---------------------------------------------------------------------------

def discover_pmids_from_subsets() -> tuple[set[str], set[str], list[str]]:
    """Load public subset manifests and extract PMIDs from referenced task files.

    Returns:
        (all_pmids, fulltext_pmids, directness_task_ids)
    """
    all_pmids: set[str] = set()
    fulltext_pmids: set[str] = set()
    directness_task_ids: list[str] = []

    for task_type, config in SUBSET_CONFIG.items():
        manifest_path = SUBSETS_DIR / config["manifest"]
        manifest = json.loads(manifest_path.read_text())
        task_ids = manifest["task_ids"]

        if task_type == "directness":
            directness_task_ids = task_ids

        for task_id in task_ids:
            task_path = config["task_dir"] / f"{task_id}.json"
            if not task_path.exists():
                print(f"  WARNING: task file not found: {task_path}")
                continue

            task = json.loads(task_path.read_text())

            if task_type == "directness":
                for study in task["studies"]:
                    pmid = study.get("pmid")
                    if pmid:
                        all_pmids.add(pmid)
            else:
                pmid = task.get("pmid")
                if pmid:
                    all_pmids.add(pmid)
                    if config["fulltext_required"]:
                        fulltext_pmids.add(pmid)

    return all_pmids, fulltext_pmids, directness_task_ids


def parse_ris_pmids(ris_path: Path) -> set[str]:
    """Extract PMIDs from the AN field of a RIS file."""
    pmids: set[str] = set()
    for line in ris_path.read_text().splitlines():
        match = re.match(r"^AN\s+-\s+(\S+)", line)
        if match:
            pmids.add(match.group(1))
    return pmids


def discover_and_validate() -> tuple[set[str], set[str], list[str]]:
    """Discover PMIDs from subsets and RIS, cross-validate, return results.

    Returns:
        (all_pmids, fulltext_pmids, directness_task_ids)
    """
    print("Discovering PMIDs from subset manifests...")
    subset_pmids, fulltext_pmids, directness_task_ids = discover_pmids_from_subsets()
    print(f"  Found {len(subset_pmids)} unique PMIDs from task files")
    print(f"  {len(fulltext_pmids)} require fulltext (RoB + eligibility)")
    print(f"  {len(subset_pmids - fulltext_pmids)} abstract only (directness)")

    print("Reading PMIDs from studies.ris...")
    ris_pmids = parse_ris_pmids(STUDIES_RIS)
    print(f"  Found {len(ris_pmids)} PMIDs in studies.ris")

    # Cross-validate
    in_subsets_only = subset_pmids - ris_pmids
    in_ris_only = ris_pmids - subset_pmids
    if in_subsets_only or in_ris_only:
        print("\nERROR: PMID mismatch between subsets and studies.ris!")
        if in_subsets_only:
            print(f"  In subsets but not in studies.ris: {sorted(in_subsets_only)}")
        if in_ris_only:
            print(f"  In studies.ris but not in subsets: {sorted(in_ris_only)}")
        sys.exit(1)

    print(f"  Cross-validation passed: {len(subset_pmids)} PMIDs match\n")
    return subset_pmids, fulltext_pmids, directness_task_ids


# ---------------------------------------------------------------------------
# Paper store helpers
# ---------------------------------------------------------------------------

def read_paper(pmid: str) -> dict | None:
    """Read a paper store entry, or None if it doesn't exist."""
    path = PAPERS_DIR / f"{pmid}.json"
    if path.exists():
        return json.loads(path.read_text())
    return None


def write_paper(paper: dict) -> None:
    """Write a paper store entry."""
    PAPERS_DIR.mkdir(parents=True, exist_ok=True)
    path = PAPERS_DIR / f"{paper['pmid']}.json"
    path.write_text(json.dumps(paper, indent=2, ensure_ascii=False) + "\n")


# ---------------------------------------------------------------------------
# download command
# ---------------------------------------------------------------------------

def cmd_download(args: argparse.Namespace) -> None:
    all_pmids, fulltext_pmids, directness_task_ids = discover_and_validate()

    # --- Step 1: Fetch abstracts from PubMed ---
    print("Step 1: Fetching abstracts from PubMed...")
    missing_abstract = []
    already_present = 0
    for pmid in sorted(all_pmids):
        paper = read_paper(pmid)
        if paper and paper.get("abstract"):
            already_present += 1
        else:
            missing_abstract.append(pmid)

    print(f"  {already_present} already in paper store, {len(missing_abstract)} to fetch")

    if missing_abstract:
        fetched = batch_fetch_abstracts_and_titles(missing_abstract)
        fetch_ok = 0
        fetch_fail = 0
        for pmid in missing_abstract:
            info = fetched.get(pmid, {})
            abstract = info.get("abstract")
            title = info.get("pubmed_title") or ""

            if not abstract:
                fetch_fail += 1
                print(f"  WARNING: no abstract returned for PMID {pmid}")
                continue

            existing = read_paper(pmid)
            if existing:
                existing["abstract"] = abstract
                existing["title"] = title or existing.get("title", "")
                write_paper(existing)
            else:
                write_paper({
                    "pmid": pmid,
                    "title": title,
                    "doi": None,
                    "abstract": abstract,
                    "fulltext": None,
                })
            fetch_ok += 1

        print(f"  Fetched {fetch_ok}, failed {fetch_fail}\n")
    else:
        print("  Nothing to fetch\n")

    # --- Step 2: Backfill directness task abstracts ---
    print("Step 2: Backfilling directness task abstracts...")
    config = SUBSET_CONFIG["directness"]
    backfilled = 0
    skipped = 0
    for task_id in directness_task_ids:
        task_path = config["task_dir"] / f"{task_id}.json"
        if not task_path.exists():
            continue

        task = json.loads(task_path.read_text())
        needs_update = False

        for study in task["studies"]:
            if study.get("abstract"):
                continue
            pmid = study.get("pmid")
            if not pmid:
                continue
            paper = read_paper(pmid)
            if not paper:
                continue
            study["title"] = paper.get("title", "")
            study["abstract"] = paper.get("abstract", "")
            needs_update = True

        if needs_update:
            task_path.write_text(json.dumps(task, indent=2, ensure_ascii=False) + "\n")
            backfilled += 1
        else:
            skipped += 1

    print(f"  Backfilled {backfilled} tasks, {skipped} already complete\n")

    # --- Step 3: Download OA PDFs ---
    if args.abstracts_only:
        print("Skipping PDF download (--abstracts-only)\n")
        print("Done. Abstracts ready for all tasks.")
        return

    print("Step 3: Downloading open-access PDFs...")
    PDFS_DIR.mkdir(parents=True, exist_ok=True)

    # Find PMIDs that need PDFs
    need_pdf = []
    pdf_already = 0
    for pmid in sorted(fulltext_pmids):
        if (PDFS_DIR / f"{pmid}.pdf").exists():
            pdf_already += 1
        else:
            need_pdf.append(pmid)

    print(f"  {pdf_already} PDFs already on disk, {len(need_pdf)} to download")

    if not need_pdf:
        print("  Nothing to download\n")
        print("Done. Run 'parse' to extract fulltext with GROBID.")
        return

    # Check OA availability
    print("  Checking OA availability via OpenAlex...")
    oa_info = check_openalex_oa(need_pdf)
    print("  Checking PMC availability...")
    pmc_info = check_pmc_availability(need_pdf)

    # Download PDFs
    downloaded = 0
    missing: list[dict] = []

    for pmid in need_pdf:
        oa = oa_info.get(pmid, {})
        pmcid = pmc_info.get(pmid)

        # Build URL list in priority order
        urls = []
        if oa.get("pdf_url"):
            urls.append(oa["pdf_url"])
        if oa.get("oa_url"):
            urls.append(oa["oa_url"])
        if pmcid:
            urls.append(
                f"https://www.ncbi.nlm.nih.gov/pmc/articles/{pmcid}/pdf/"
            )

        if not urls:
            missing.append({"pmid": pmid, "reason": "not_oa"})
            continue

        pdf_ok = False
        for url in urls:
            try:
                time.sleep(API_DELAY)
                resp = requests.get(
                    url, timeout=30, headers={"User-Agent": "SynthesisBench/0.1"}
                )
                resp.raise_for_status()

                if not resp.content[:5].startswith(b"%PDF"):
                    continue  # try next URL

                (PDFS_DIR / f"{pmid}.pdf").write_bytes(resp.content)
                downloaded += 1
                pdf_ok = True
                print(f"  [{downloaded}] {pmid}: OK")
                break
            except requests.RequestException:
                continue  # try next URL

        if not pdf_ok:
            missing.append({
                "pmid": pmid,
                "reason": "download_failed",
                "urls_tried": urls,
            })

    # Write missing.json
    missing_path = PDFS_DIR / "missing.json"
    missing_path.write_text(json.dumps(missing, indent=2) + "\n")

    total_needed = len(fulltext_pmids)
    total_have = pdf_already + downloaded
    print(f"\n  Downloaded {downloaded}/{len(need_pdf)} PDFs")
    print(f"  Total coverage: {total_have}/{total_needed} fulltext-required papers")
    if missing:
        print(f"  {len(missing)} unavailable — see {missing_path}")
        print("  Place PDFs manually as data/pdfs/{pmid}.pdf and re-run 'parse'")

    print("\nDone. Run 'parse' to extract fulltext with GROBID.")


# ---------------------------------------------------------------------------
# parse command
# ---------------------------------------------------------------------------

def cmd_parse(args: argparse.Namespace) -> None:
    all_pmids, fulltext_pmids, _ = discover_and_validate()

    grobid_url = args.grobid_url

    # Find PMIDs with PDF but no fulltext
    need_parse = []
    already_parsed = 0
    no_pdf = 0
    for pmid in sorted(fulltext_pmids):
        pdf_path = PDFS_DIR / f"{pmid}.pdf"
        if not pdf_path.exists():
            no_pdf += 1
            continue
        paper = read_paper(pmid)
        if paper and paper.get("fulltext"):
            already_parsed += 1
            continue
        need_parse.append(pmid)

    print(f"  {already_parsed} already have fulltext")
    print(f"  {no_pdf} missing PDF (download first or place manually)")
    print(f"  {len(need_parse)} to parse\n")

    if not need_parse:
        total = already_parsed + len(need_parse)
        print(f"Nothing to parse. Coverage: {already_parsed}/{len(fulltext_pmids)} fulltext-required papers.")
        return

    # Check GROBID
    if not check_grobid_available(grobid_url):
        print(f"ERROR: GROBID not available at {grobid_url}")
        print()
        print("Start GROBID with Docker:")
        print("  docker run -d --name grobid -p 8070:8070 lfoppiano/grobid:0.8.0")
        print()
        print("Then re-run: python setup_papers.py parse")
        sys.exit(1)

    print(f"GROBID available at {grobid_url}")
    parser = GrobidSectionParser(grobid_url)

    parsed = 0
    failed = 0
    for i, pmid in enumerate(need_parse):
        pdf_path = PDFS_DIR / f"{pmid}.pdf"
        try:
            sections = parser.parse_pdf(str(pdf_path))
            paper = read_paper(pmid)
            if not paper:
                paper = {
                    "pmid": pmid,
                    "title": "",
                    "doi": None,
                    "abstract": None,
                    "fulltext": None,
                }
            paper["fulltext"] = {"sections": sections, "source": "grobid"}
            write_paper(paper)
            parsed += 1
            print(f"  [{i + 1}/{len(need_parse)}] {pmid}: {len(sections)} sections")
        except GrobidParsingError as e:
            failed += 1
            print(f"  [{i + 1}/{len(need_parse)}] {pmid}: FAILED ({e})")

    print(f"\nParsed {parsed}/{len(need_parse)}")
    if failed:
        print(f"{failed} failed")
    total_with_fulltext = already_parsed + parsed
    print(f"Coverage: {total_with_fulltext}/{len(fulltext_pmids)} fulltext-required papers")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Set up the paper store for EvidenceSynthBench.",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # download
    dl = subparsers.add_parser(
        "download",
        help="Download abstracts and OA PDFs",
    )
    dl.add_argument(
        "--abstracts-only",
        action="store_true",
        help="Only fetch abstracts, skip PDF download",
    )

    # parse
    ps = subparsers.add_parser(
        "parse",
        help="Parse downloaded PDFs with GROBID",
    )
    ps.add_argument(
        "--grobid-url",
        default="http://localhost:8070",
        help="GROBID server URL (default: http://localhost:8070)",
    )

    args = parser.parse_args()

    if args.command == "download":
        cmd_download(args)
    elif args.command == "parse":
        cmd_parse(args)


if __name__ == "__main__":
    main()