Deep Research

#!/usr/bin/env python3
"""Generate and manage BibTeX entries from JSONL paper records.

Citation keys: firstAuthorLastNameYearFirstTitleWord (e.g., vaswani2017attention).

Usage:
    python bibtex_manager.py --jsonl paper_db.jsonl --output references.bib
    python bibtex_manager.py --jsonl paper_db.jsonl  # stdout
"""

import argparse
import json
import re
import sys
import unicodedata


def normalize_name(name: str) -> str:
    """Normalize a name: strip accents, lowercase."""
    nfkd = unicodedata.normalize("NFKD", name)
    ascii_name = nfkd.encode("ascii", "ignore").decode("ascii")
    return ascii_name.strip()


def last_name(author: str) -> str:
    """Extract last name from an author string."""
    author = normalize_name(author)
    parts = author.split()
    if not parts:
        return "unknown"
    return parts[-1].lower()


def make_citation_key(paper: dict) -> str:
    """Generate a citation key: lastNameYearWord."""
    authors = paper.get("authors", [])
    first_author = last_name(authors[0]) if authors else "unknown"
    # Remove non-alphanumeric from author
    first_author = re.sub(r"[^a-z]", "", first_author)

    year = paper.get("year", "")
    if not year:
        pub = paper.get("published", "") or paper.get("publicationDate", "") or ""
        if pub and len(pub) >= 4:
            year = pub[:4]
        else:
            year = "nd"

    title = paper.get("title", "")
    skip_words = {"a", "an", "the", "on", "in", "of", "for", "to", "with", "and", "or"}
    title_words = re.findall(r"[a-z]+", title.lower())
    first_word = "paper"
    for w in title_words:
        if w not in skip_words and len(w) > 2:
            first_word = w
            break

    return f"{first_author}{year}{first_word}"


def escape_bibtex(text: str) -> str:
    """Escape special BibTeX characters."""
    text = text.replace("&", r"\&")
    text = text.replace("%", r"\%")
    text = text.replace("#", r"\#")
    text = text.replace("_", r"\_")
    return text


def format_authors_bibtex(authors: list[str]) -> str:
    """Format authors for BibTeX (Name1 and Name2 and Name3)."""
    if not authors:
        return "Unknown"
    return " and ".join(authors)


def paper_to_bibtex(paper: dict, key: str) -> str:
    """Convert a paper record to a BibTeX entry."""
    arxiv_id = paper.get("arxiv_id", "")
    venue = paper.get("venue", "")
    year = paper.get("year", "")
    if not year:
        pub = paper.get("published", "") or paper.get("publicationDate", "") or ""
        year = pub[:4] if len(pub) >= 4 else ""

    authors = format_authors_bibtex(paper.get("authors", []))
    title = paper.get("title", "")
    abstract_text = paper.get("abstract", "")

    # Determine entry type
    conf_keywords = ["conference", "proceedings", "icml", "neurips", "iclr", "acl", "emnlp", "cvpr", "aaai"]
    journal_keywords = ["journal", "transactions", "review"]

    if venue and any(kw in venue.lower() for kw in conf_keywords):
        entry_type = "inproceedings"
    elif venue and any(kw in venue.lower() for kw in journal_keywords):
        entry_type = "article"
    elif arxiv_id:
        entry_type = "article"
    else:
        entry_type = "misc"

    lines = [f"@{entry_type}{{{key},"]
    lines.append(f"  title = {{{escape_bibtex(title)}}},")
    lines.append(f"  author = {{{authors}}},")
    if year:
        lines.append(f"  year = {{{year}}},")
    if venue:
        if entry_type == "inproceedings":
            lines.append(f"  booktitle = {{{escape_bibtex(venue)}}},")
        elif entry_type == "article" and not arxiv_id:
            lines.append(f"  journal = {{{escape_bibtex(venue)}}},")
    if arxiv_id:
        lines.append(f"  eprint = {{{arxiv_id}}},")
        lines.append(f"  archivePrefix = {{arXiv}},")
        if not venue:
            lines.append(f"  journal = {{arXiv preprint arXiv:{arxiv_id}}},")
    url = paper.get("url", "")
    if url:
        lines.append(f"  url = {{{url}}},")
    if abstract_text:
        short_abstract = abstract_text[:500]
        lines.append(f"  abstract = {{{escape_bibtex(short_abstract)}}},")
    lines.append("}")

    return "\n".join(lines)


def load_jsonl(path: str) -> list[dict]:
    """Load records from a JSONL file."""
    records = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def generate_bibtex(records: list[dict]) -> str:
    """Generate BibTeX for all records, with deduplication by key."""
    entries = {}
    for rec in records:
        key = make_citation_key(rec)
        original_key = key
        suffix_idx = 0
        while key in entries:
            suffix_idx += 1
            key = f"{original_key}{chr(96 + suffix_idx)}"  # a, b, c...
        entries[key] = paper_to_bibtex(rec, key)

    return "\n\n".join(entries.values()) + "\n"


def main():
    parser = argparse.ArgumentParser(description="Generate BibTeX from JSONL paper records")
    parser.add_argument("--jsonl", required=True, help="Input JSONL file with paper records")
    parser.add_argument("--output", "-o", help="Output .bib file (default: stdout)")
    parser.add_argument("--keys-only", action="store_true", help="Only print citation keys")
    args = parser.parse_args()

    records = load_jsonl(args.jsonl)

    if args.keys_only:
        seen = set()
        for rec in records:
            key = make_citation_key(rec)
            original_key = key
            suffix_idx = 0
            while key in seen:
                suffix_idx += 1
                key = f"{original_key}{chr(96 + suffix_idx)}"
            seen.add(key)
            title = rec.get("title", "")[:60]
            print(f"{key}\t{title}")
        return

    bibtex = generate_bibtex(records)

    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(bibtex)
        print(f"Written {len(records)} entries to {args.output}", file=sys.stderr)
    else:
        print(bibtex)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Compile research notes into a final report with numbered citations and BibTeX.

Reads notes/*.md + paper_db.jsonl + code_repos.md from a topic directory
and generates report.md with proper [1], [2] citations and references.bib.

Self-contained: uses only stdlib.

Usage:
    python compile_report.py --topic-dir output/long-horizon-reasoning/
"""

import argparse
import json
import os
import re
import sys
from collections import Counter


def load_papers(jsonl_path: str) -> list[dict]:
    """Load paper records from JSONL."""
    papers = []
    if not os.path.exists(jsonl_path):
        return papers
    with open(jsonl_path) as f:
        for line in f:
            line = line.strip()
            if line:
                papers.append(json.loads(line))
    return papers


def load_text(path: str) -> str:
    """Load a text file, return empty string if missing."""
    if not os.path.exists(path):
        return ""
    with open(path) as f:
        return f.read()


def make_cite_key(paper: dict) -> str:
    """Generate a citation key from a paper record."""
    arxiv_id = paper.get("arxiv_id", "")
    if arxiv_id:
        return arxiv_id.replace("/", "_").replace(".", "_")
    # Fallback: first author last name + year
    authors = paper.get("authors", [])
    first = authors[0].split()[-1].lower() if authors else "unknown"
    year = paper.get("year") or paper.get("published", "")[:4] or "0000"
    title_word = paper.get("title", "").split()[0].lower() if paper.get("title") else "paper"
    return f"{first}{year}_{title_word}"


def paper_to_bibtex(paper: dict, cite_key: str) -> str:
    """Convert a paper record to BibTeX entry."""
    authors = " and ".join(paper.get("authors", ["Unknown"]))
    title = paper.get("title", "Unknown")
    year = paper.get("year") or paper.get("published", "")[:4] or ""
    venue = paper.get("venue", "")
    arxiv_id = paper.get("arxiv_id", "")
    url = paper.get("url", "")

    if venue:
        entry_type = "inproceedings"
        venue_field = f"  booktitle = {{{venue}}},"
    elif arxiv_id:
        entry_type = "article"
        venue_field = f"  journal = {{arXiv preprint arXiv:{arxiv_id}}},"
    else:
        entry_type = "article"
        venue_field = ""

    lines = [
        f"@{entry_type}{{{cite_key},",
        f"  title = {{{title}}},",
        f"  author = {{{authors}}},",
        f"  year = {{{year}}},",
    ]
    if venue_field:
        lines.append(venue_field)
    if url:
        lines.append(f"  url = {{{url}}},")
    lines.append("}")
    return "\n".join(lines)


def build_citation_map(papers: list[dict]) -> dict[str, tuple[int, dict]]:
    """Build a map from various paper identifiers to (number, paper).

    Keys used: arxiv_id, paperId, cite_key, and partial title matches.
    """
    cite_map = {}
    for i, paper in enumerate(papers, 1):
        cite_key = make_cite_key(paper)
        paper["_cite_key"] = cite_key
        paper["_cite_num"] = i

        if paper.get("arxiv_id"):
            cite_map[paper["arxiv_id"]] = (i, paper)
        if paper.get("paperId"):
            cite_map[paper["paperId"]] = (i, paper)
        cite_map[cite_key] = (i, paper)

    return cite_map


def replace_citations(text: str, cite_map: dict[str, tuple[int, dict]]) -> str:
    """Replace [@key] citations in text with numbered [N] references."""
    def replacer(match):
        key = match.group(1)
        if key in cite_map:
            num, _ = cite_map[key]
            return f"[{num}]"
        return match.group(0)  # leave unchanged if not found

    return re.sub(r"\[@([^\]]+)\]", replacer, text)


def compute_stats(papers: list[dict]) -> str:
    """Generate summary statistics section."""
    if not papers:
        return "No papers in database."

    lines = []
    lines.append(f"**Total papers**: {len(papers)}")

    # By year (normalize to int for consistent grouping)
    def normalize_year(y):
        if y is None:
            return "Unknown"
        try:
            return int(y)
        except (ValueError, TypeError):
            return "Unknown"

    years = Counter(normalize_year(p.get("year")) for p in papers)
    lines.append("\n**Papers by year**:")
    for year in sorted(years.keys(), key=lambda x: (0, 0) if x == "Unknown" else (1, x), reverse=True):
        lines.append(f"- {year}: {years[year]}")

    # By venue (top 10)
    venues = Counter(p.get("venue", "") or "Preprint" for p in papers)
    top_venues = venues.most_common(10)
    lines.append("\n**Top venues**:")
    for venue, count in top_venues:
        lines.append(f"- {venue}: {count}")

    # Top cited
    cited = sorted(papers, key=lambda p: p.get("citationCount", 0) or 0, reverse=True)[:10]
    if any(p.get("citationCount", 0) for p in cited):
        lines.append("\n**Most cited papers**:")
        for p in cited:
            cc = p.get("citationCount", 0) or 0
            if cc > 0:
                lines.append(f"- [{p.get('_cite_num', '?')}] {p['title'][:80]} ({cc} citations)")

    return "\n".join(lines)


def load_notes(topic_dir: str) -> str:
    """Load note files from phase-based directory structure.

    Searches for notes in this order:
    1. Phase-based: phase1_frontier/frontier.md, phase2_survey/survey.md, etc.
    2. Legacy notes/ directory: notes/frontier.md, notes/survey.md, etc.
    3. Legacy single file: notes.md
    """
    # Phase-based note locations (new structure)
    phase_notes = [
        ("phase1_frontier", "frontier.md"),
        ("phase2_survey", "survey.md"),
        ("phase3_deep_dive", "deep_dive.md"),
        ("phase5_synthesis", "synthesis.md"),
        ("phase5_synthesis", "gaps.md"),
    ]
    parts = []
    for subdir, filename in phase_notes:
        path = os.path.join(topic_dir, subdir, filename)
        if os.path.exists(path):
            content = load_text(path)
            if content.strip():
                parts.append(content)

    if parts:
        return "\n\n---\n\n".join(parts)

    # Fallback: legacy notes/ directory
    notes_dir = os.path.join(topic_dir, "notes")
    legacy_files = ["frontier.md", "survey.md", "deep_dive.md", "synthesis.md", "gaps.md"]
    if os.path.isdir(notes_dir):
        for name in legacy_files:
            path = os.path.join(notes_dir, name)
            if os.path.exists(path):
                content = load_text(path)
                if content.strip():
                    parts.append(content)

    if not parts:
        fallback = load_text(os.path.join(topic_dir, "notes.md"))
        if fallback.strip():
            parts.append(fallback)

    return "\n\n---\n\n".join(parts)


def compile_report(topic_dir: str):
    """Compile all materials into a final report."""
    paper_db_path = os.path.join(topic_dir, "paper_db.jsonl")

    # Code repos: phase-based first, then legacy fallbacks
    code_path = os.path.join(topic_dir, "phase4_code", "code_repos.md")
    if not os.path.exists(code_path):
        code_path = os.path.join(topic_dir, "code_repos.md")
    if not os.path.exists(code_path):
        code_path = os.path.join(topic_dir, "code_resources.md")

    # Report output: phase-based directory
    report_dir = os.path.join(topic_dir, "phase6_report")
    os.makedirs(report_dir, exist_ok=True)
    report_path = os.path.join(report_dir, "report.md")
    bib_path = os.path.join(report_dir, "references.bib")

    papers = load_papers(paper_db_path)
    notes = load_notes(topic_dir)
    code_resources = load_text(code_path)

    if not papers and not notes:
        print("Warning: no papers or notes found", file=sys.stderr)

    cite_map = build_citation_map(papers)

    # Process notes: replace [@key] with [N]
    processed_notes = replace_citations(notes, cite_map)

    # Build report
    report_parts = []

    # Title
    topic_name = os.path.basename(topic_dir.rstrip("/")).replace("-", " ").title()
    report_parts.append(f"# Research Report: {topic_name}\n")
    report_parts.append(f"*Generated from {len(papers)} papers*\n")

    # Statistics
    report_parts.append("## Paper Statistics\n")
    report_parts.append(compute_stats(papers))

    # Notes (the main content)
    if processed_notes:
        report_parts.append("\n---\n")
        report_parts.append(processed_notes)

    # Code resources
    if code_resources:
        report_parts.append("\n---\n")
        report_parts.append("## Code & Tools\n")
        report_parts.append(code_resources)

    # References
    report_parts.append("\n---\n")
    report_parts.append("## References\n")
    for paper in papers:
        num = paper.get("_cite_num", "?")
        title = paper.get("title", "Unknown")
        authors = paper.get("authors", [])
        year = paper.get("year") or paper.get("published", "")[:4] or ""
        author_str = ", ".join(authors[:3])
        if len(authors) > 3:
            author_str += " et al."
        venue = paper.get("venue", "")
        venue_str = f" {venue}." if venue else ""
        url = paper.get("url", "")
        url_str = f" {url}" if url else ""
        report_parts.append(f"[{num}] {author_str}. \"{title}\". {year}.{venue_str}{url_str}\n")

    report_text = "\n".join(report_parts)

    # Write report
    with open(report_path, "w") as f:
        f.write(report_text)
    print(f"Report written to {report_path} ({len(report_text)} chars)", file=sys.stderr)

    # Write BibTeX
    bib_entries = []
    for paper in papers:
        cite_key = paper.get("_cite_key", make_cite_key(paper))
        bib_entries.append(paper_to_bibtex(paper, cite_key))

    with open(bib_path, "w") as f:
        f.write("\n\n".join(bib_entries) + "\n")
    print(f"BibTeX written to {bib_path} ({len(bib_entries)} entries)", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="Compile research notes into a report")
    parser.add_argument("--topic-dir", required=True, help="Topic output directory")
    args = parser.parse_args()

    if not os.path.isdir(args.topic_dir):
        print(f"Error: {args.topic_dir} is not a directory", file=sys.stderr)
        sys.exit(1)

    compile_report(args.topic_dir)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Download PDFs from a JSONL paper database.

Self-contained: uses only stdlib (urllib).

Features:
- Atomic downloads (.part file -> rename on success)
- PDF validation (checks %PDF header + %%EOF trailer)
- Respects rate limits (configurable delay)
- Skips already-downloaded papers

Usage:
    python download_papers.py --jsonl paper_db.jsonl --output-dir papers/
    python download_papers.py --jsonl paper_db.jsonl --output-dir papers/ --max-downloads 20 --delay 2.0
"""

import argparse
import json
import os
import sys
import time
import urllib.request


def sanitize_filename(arxiv_id: str, paper_id: str) -> str:
    """Create a safe filename from paper IDs."""
    name = arxiv_id or paper_id
    # Replace path separators and problematic chars
    name = name.replace("/", "_").replace("\\", "_").replace(":", "_")
    if not name.endswith(".pdf"):
        name += ".pdf"
    return name


def validate_pdf(path: str) -> bool:
    """Check that a file looks like a valid PDF."""
    try:
        with open(path, "rb") as f:
            header = f.read(5)
            if header != b"%PDF-":
                return False
            # Check for EOF marker in last 1KB
            f.seek(0, 2)
            size = f.tell()
            f.seek(max(0, size - 1024))
            tail = f.read()
            return b"%%EOF" in tail
    except Exception:
        return False


def download_pdf(url: str, dest: str, timeout: int = 60) -> bool:
    """Download a PDF with atomic write. Returns True on success."""
    part_path = dest + ".part"

    headers = {
        "User-Agent": "deep-research/1.0 (academic research tool)",
    }
    req = urllib.request.Request(url, headers=headers)

    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            content_type = resp.headers.get("Content-Type", "")
            # Some servers redirect to HTML (captcha, etc.)
            if "text/html" in content_type and "pdf" not in content_type:
                print(f"  Warning: got HTML instead of PDF from {url}", file=sys.stderr)
                return False

            with open(part_path, "wb") as f:
                while True:
                    chunk = resp.read(8192)
                    if not chunk:
                        break
                    f.write(chunk)

        if not validate_pdf(part_path):
            print(f"  Warning: invalid PDF from {url}", file=sys.stderr)
            os.remove(part_path)
            return False

        os.rename(part_path, dest)
        return True

    except Exception as e:
        print(f"  Error downloading {url}: {e}", file=sys.stderr)
        if os.path.exists(part_path):
            os.remove(part_path)
        return False


def load_papers(jsonl_path: str) -> list[dict]:
    """Load papers from a JSONL file."""
    papers = []
    with open(jsonl_path) as f:
        for line in f:
            line = line.strip()
            if line:
                papers.append(json.loads(line))
    return papers


def main():
    parser = argparse.ArgumentParser(description="Download PDFs from JSONL paper database")
    parser.add_argument("--jsonl", required=True, help="JSONL file with paper records")
    parser.add_argument("--output-dir", required=True, help="Directory to save PDFs")
    parser.add_argument("--max-downloads", type=int, default=50, help="Max PDFs to download")
    parser.add_argument("--delay", type=float, default=1.0, help="Seconds between downloads")
    parser.add_argument("--timeout", type=int, default=60, help="Download timeout in seconds")
    parser.add_argument("--sort-by-citations", action="store_true", help="Download most-cited first")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    papers = load_papers(args.jsonl)

    if args.sort_by_citations:
        papers.sort(key=lambda p: p.get("citationCount", 0) or 0, reverse=True)

    downloaded = 0
    skipped = 0
    failed = 0

    for paper in papers:
        if downloaded >= args.max_downloads:
            break

        pdf_url = paper.get("pdf_url", "")
        if not pdf_url:
            continue

        arxiv_id = paper.get("arxiv_id", "")
        paper_id = paper.get("paperId", "")
        filename = sanitize_filename(arxiv_id, paper_id)
        dest = os.path.join(args.output_dir, filename)

        if os.path.exists(dest):
            skipped += 1
            continue

        title = paper.get("title", "unknown")[:60]
        print(f"[{downloaded + 1}/{args.max_downloads}] {title}...", file=sys.stderr)

        if download_pdf(pdf_url, dest, timeout=args.timeout):
            size_mb = os.path.getsize(dest) / (1024 * 1024)
            print(f"  OK ({size_mb:.1f} MB)", file=sys.stderr)
            downloaded += 1
        else:
            failed += 1

        if downloaded < args.max_downloads:
            time.sleep(args.delay)

    print(f"\nDone: {downloaded} downloaded, {skipped} skipped, {failed} failed", file=sys.stderr)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Extract text from PDFs using PyMuPDF (fitz).

Features:
- Full text extraction with layout preservation
- Section detection (Abstract, Introduction, Methods, etc.)
- Batch mode for entire directories
- BibTeX-style reference extraction

Usage:
    python extract_pdf.py --pdf paper.pdf
    python extract_pdf.py --pdf-dir papers/ --output-dir texts/
    python extract_pdf.py --pdf paper.pdf --sections-only
"""

import argparse
import os
import re
import sys

try:
    import fitz  # PyMuPDF
except ImportError:
    print("Error: PyMuPDF not installed. Run: pip install PyMuPDF", file=sys.stderr)
    sys.exit(1)


SECTION_PATTERNS = [
    (r"^\s*abstract\s*$", "Abstract"),
    (r"^\s*\d*\.?\s*introduction\s*$", "Introduction"),
    (r"^\s*\d*\.?\s*related\s+work", "Related Work"),
    (r"^\s*\d*\.?\s*background", "Background"),
    (r"^\s*\d*\.?\s*method(?:s|ology)?", "Methods"),
    (r"^\s*\d*\.?\s*(?:proposed\s+)?(?:approach|framework|model|system)", "Methods"),
    (r"^\s*\d*\.?\s*experiment(?:s|al)?(?:\s+(?:setup|results))?", "Experiments"),
    (r"^\s*\d*\.?\s*results?(?:\s+and\s+(?:discussion|analysis))?", "Results"),
    (r"^\s*\d*\.?\s*evaluation", "Evaluation"),
    (r"^\s*\d*\.?\s*discussion", "Discussion"),
    (r"^\s*\d*\.?\s*(?:conclusion|concluding)", "Conclusion"),
    (r"^\s*\d*\.?\s*limitation", "Limitations"),
    (r"^\s*\d*\.?\s*(?:future\s+work|outlook)", "Future Work"),
    (r"^\s*\d*\.?\s*(?:acknowledge?ment)", "Acknowledgements"),
    (r"^\s*\d*\.?\s*references?\s*$", "References"),
    (r"^\s*\d*\.?\s*(?:appendix|supplementary)", "Appendix"),
]


def extract_text(pdf_path: str) -> str:
    """Extract full text from a PDF file."""
    doc = fitz.open(pdf_path)
    pages = []
    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text("text")
        if text.strip():
            pages.append(text)
    doc.close()
    return "\n\n".join(pages)


def detect_sections(text: str) -> list[tuple[str, str]]:
    """Detect sections in extracted text. Returns list of (section_name, content)."""
    lines = text.split("\n")
    sections = []
    current_section = "Preamble"
    current_lines = []

    for line in lines:
        stripped = line.strip()
        if not stripped:
            current_lines.append("")
            continue

        matched = False
        # Check if this line is a section header
        # Heuristic: short line (< 80 chars) matching a known pattern
        if len(stripped) < 80:
            for pattern, section_name in SECTION_PATTERNS:
                if re.match(pattern, stripped, re.IGNORECASE):
                    # Save previous section
                    if current_lines:
                        content = "\n".join(current_lines).strip()
                        if content:
                            sections.append((current_section, content))
                    current_section = section_name
                    current_lines = []
                    matched = True
                    break

        if not matched:
            current_lines.append(line)

    # Don't forget the last section
    if current_lines:
        content = "\n".join(current_lines).strip()
        if content:
            sections.append((current_section, content))

    return sections


def extract_with_sections(pdf_path: str) -> dict:
    """Extract text and identify sections."""
    text = extract_text(pdf_path)
    sections = detect_sections(text)
    return {
        "full_text": text,
        "sections": {name: content for name, content in sections},
        "section_order": [name for name, _ in sections],
    }


def format_sections(result: dict, sections_only: bool = False) -> str:
    """Format extracted text with section markers."""
    if sections_only:
        output = []
        for name in result["section_order"]:
            content = result["sections"].get(name, "")
            output.append(f"## {name}\n\n{content}")
        return "\n\n---\n\n".join(output)
    else:
        return result["full_text"]


def process_directory(pdf_dir: str, output_dir: str, sections_only: bool = False):
    """Process all PDFs in a directory."""
    os.makedirs(output_dir, exist_ok=True)
    pdf_files = sorted(f for f in os.listdir(pdf_dir) if f.lower().endswith(".pdf"))

    for i, filename in enumerate(pdf_files, 1):
        pdf_path = os.path.join(pdf_dir, filename)
        txt_name = os.path.splitext(filename)[0] + ".txt"
        txt_path = os.path.join(output_dir, txt_name)

        print(f"[{i}/{len(pdf_files)}] {filename}...", file=sys.stderr)

        try:
            result = extract_with_sections(pdf_path)
            text = format_sections(result, sections_only)
            with open(txt_path, "w") as f:
                f.write(text)
            print(f"  OK ({len(text)} chars)", file=sys.stderr)
        except Exception as e:
            print(f"  Error: {e}", file=sys.stderr)

    print(f"\nProcessed {len(pdf_files)} PDFs", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="Extract text from PDFs using PyMuPDF")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--pdf", help="Single PDF file to extract")
    group.add_argument("--pdf-dir", help="Directory of PDFs to process")
    parser.add_argument("--output-dir", help="Output directory for batch mode")
    parser.add_argument("--sections-only", action="store_true", help="Output only detected sections")
    args = parser.parse_args()

    if args.pdf:
        result = extract_with_sections(args.pdf)
        print(format_sections(result, args.sections_only))
    elif args.pdf_dir:
        if not args.output_dir:
            print("Error: --output-dir required with --pdf-dir", file=sys.stderr)
            sys.exit(1)
        process_directory(args.pdf_dir, args.output_dir, args.sections_only)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""JSONL paper database management.

Subcommands: add, search, merge, tag, stats, export.
Deduplication by title similarity (Jaccard on word tokens, threshold 0.8).

Usage:
    python paper_db.py merge --inputs arxiv.jsonl s2.jsonl --output merged.jsonl
    python paper_db.py stats --input paper_db.jsonl
    python paper_db.py search --input paper_db.jsonl --query "transformer"
    python paper_db.py tag --input paper_db.jsonl --ids "2401.12345" --tags core method
    python paper_db.py add --input paper_db.jsonl --record '{"title":"...", "arxiv_id":"..."}'
    python paper_db.py export --input paper_db.jsonl --format csv
"""

import argparse
import csv
import io
import json
import os
import re
import sys


def tokenize(text: str) -> set[str]:
    """Tokenize text into lowercase word set."""
    return set(re.findall(r"[a-z0-9]+", text.lower()))


def jaccard(a: set, b: set) -> float:
    """Jaccard similarity between two sets."""
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)


def load_jsonl(path: str) -> list[dict]:
    """Load records from a JSONL file."""
    records = []
    if not os.path.exists(path):
        return records
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def save_jsonl(records: list[dict], path: str):
    """Save records to a JSONL file."""
    with open(path, "w", encoding="utf-8") as f:
        for rec in records:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def get_paper_id(paper: dict) -> str:
    """Get the canonical ID for a paper."""
    return paper.get("arxiv_id") or paper.get("paperId") or ""


def deduplicate(records: list[dict], threshold: float = 0.8) -> list[dict]:
    """Remove duplicate papers by title similarity."""
    seen_titles: list[set[str]] = []
    seen_ids: set[str] = set()
    unique = []

    for rec in records:
        pid = get_paper_id(rec)
        title = rec.get("title", "")

        # Exact ID match
        if pid and pid in seen_ids:
            continue

        # Title similarity check
        title_tokens = tokenize(title)
        is_dup = False
        for prev_tokens in seen_titles:
            if jaccard(title_tokens, prev_tokens) >= threshold:
                is_dup = True
                break

        if is_dup:
            continue

        if pid:
            seen_ids.add(pid)
        seen_titles.append(title_tokens)
        unique.append(rec)

    return unique


def merge_databases(inputs: list[str], output: str, threshold: float = 0.8):
    """Merge multiple JSONL files with deduplication."""
    all_records = []
    for path in inputs:
        records = load_jsonl(path)
        print(f"Loaded {len(records)} from {path}", file=sys.stderr)
        all_records.extend(records)

    merged = deduplicate(all_records, threshold)
    save_jsonl(merged, output)
    print(f"Merged: {len(all_records)} -> {len(merged)} unique papers -> {output}", file=sys.stderr)


def filter_db(db_path: str, output: str, *, min_score: float = 0.0, max_papers: int = 0,
              require_keywords: list[str] | None = None):
    """Filter papers by affinity_score threshold and optional keyword relevance."""
    records = load_jsonl(db_path)
    kept = []
    for rec in records:
        score = rec.get("affinity_score")
        if score is not None and score < min_score:
            continue
        if score is None and require_keywords:
            title_lower = rec.get("title", "").lower()
            if not any(k in title_lower for k in require_keywords):
                continue
        kept.append(rec)

    # Sort by score descending (None at end)
    kept.sort(key=lambda r: -(r.get("affinity_score") or 0))

    if max_papers > 0 and len(kept) > max_papers:
        kept = kept[:max_papers]

    save_jsonl(kept, output)
    print(f"Filtered: {len(records)} -> {len(kept)} papers -> {output}", file=sys.stderr)


def search_db(db_path: str, query: str, field: str = "title") -> list[dict]:
    """Search papers by keyword match in a field."""
    records = load_jsonl(db_path)
    query_lower = query.lower()
    results = []

    for rec in records:
        value = rec.get(field, "")
        if isinstance(value, list):
            value = " ".join(str(v) for v in value)
        if query_lower in str(value).lower():
            results.append(rec)

    return results


def tag_papers(db_path: str, ids: list[str], tags: list[str]):
    """Add tags to specific papers."""
    records = load_jsonl(db_path)
    tagged = 0

    for rec in records:
        pid = get_paper_id(rec)
        if pid in ids:
            existing = rec.get("tags", [])
            rec["tags"] = sorted(set(existing + tags))
            tagged += 1

    save_jsonl(records, db_path)
    print(f"Tagged {tagged} papers with {tags}", file=sys.stderr)


def compute_stats(db_path: str) -> dict:
    """Compute statistics about the paper database."""
    records = load_jsonl(db_path)
    if not records:
        return {"total": 0}

    sources = {}
    years = {}
    venues = {}
    with_abstract = 0
    with_pdf = 0
    total_citations = 0
    tags_dist = {}
    peer_reviewed_count = 0

    for rec in records:
        src = rec.get("source", "unknown")
        sources[src] = sources.get(src, 0) + 1

        year = rec.get("year")
        if year:
            years[year] = years.get(year, 0) + 1

        venue = rec.get("venue", "")
        if venue:
            venues[venue] = venues.get(venue, 0) + 1

        if rec.get("abstract"):
            with_abstract += 1
        if rec.get("pdf_url") or rec.get("pdf_path"):
            with_pdf += 1
        if rec.get("peer_reviewed"):
            peer_reviewed_count += 1

        total_citations += rec.get("citationCount", 0) or 0

        for tag in rec.get("tags", []):
            tags_dist[tag] = tags_dist.get(tag, 0) + 1

    return {
        "total": len(records),
        "peer_reviewed": peer_reviewed_count,
        "preprint_only": len(records) - peer_reviewed_count,
        "sources": sources,
        "years": dict(sorted(years.items(), key=lambda x: str(x[0]))),
        "top_venues": dict(sorted(venues.items(), key=lambda x: -x[1])[:10]),
        "with_abstract": with_abstract,
        "with_pdf": with_pdf,
        "total_citations": total_citations,
        "avg_citations": round(total_citations / len(records), 1),
        "tags": tags_dist,
    }


def export_csv(db_path: str) -> str:
    """Export paper DB as CSV."""
    records = load_jsonl(db_path)
    if not records:
        return ""

    fields = ["arxiv_id", "paperId", "title", "authors", "year", "venue",
              "citationCount", "pdf_url", "tags", "source"]

    output = io.StringIO()
    writer = csv.DictWriter(output, fieldnames=fields, extrasaction="ignore")
    writer.writeheader()
    for rec in records:
        row = dict(rec)
        if isinstance(row.get("authors"), list):
            row["authors"] = "; ".join(row["authors"])
        if isinstance(row.get("tags"), list):
            row["tags"] = "; ".join(row["tags"])
        writer.writerow(row)

    return output.getvalue()


def main():
    parser = argparse.ArgumentParser(description="JSONL paper database management")
    sub = parser.add_subparsers(dest="command", required=True)

    # merge
    p_merge = sub.add_parser("merge", help="Merge multiple JSONL files with deduplication")
    p_merge.add_argument("--inputs", nargs="+", required=True, help="Input JSONL files")
    p_merge.add_argument("--output", required=True, help="Output JSONL file")
    p_merge.add_argument("--threshold", type=float, default=0.8, help="Title similarity threshold")

    # search
    p_search = sub.add_parser("search", help="Search papers by keyword")
    p_search.add_argument("--input", required=True, help="Paper DB JSONL file")
    p_search.add_argument("--query", required=True, help="Search query")
    p_search.add_argument("--field", default="title", help="Field to search (default: title)")

    # tag
    p_tag = sub.add_parser("tag", help="Add tags to papers")
    p_tag.add_argument("--input", required=True, help="Paper DB JSONL file")
    p_tag.add_argument("--ids", nargs="+", required=True, help="Paper IDs to tag")
    p_tag.add_argument("--tags", nargs="+", required=True, help="Tags to add")

    # add
    p_add = sub.add_parser("add", help="Add a paper record")
    p_add.add_argument("--input", required=True, help="Paper DB JSONL file")
    p_add.add_argument("--record", required=True, help="JSON string of paper record")

    # stats
    p_stats = sub.add_parser("stats", help="Show database statistics")
    p_stats.add_argument("--input", required=True, help="Paper DB JSONL file")

    # filter
    p_filter = sub.add_parser("filter", help="Filter papers by score/keywords")
    p_filter.add_argument("--input", required=True, help="Paper DB JSONL file")
    p_filter.add_argument("--output", "-o", required=True, help="Output JSONL file")
    p_filter.add_argument("--min-score", type=float, default=0.0, help="Minimum affinity_score threshold")
    p_filter.add_argument("--max-papers", type=int, default=0, help="Maximum number of papers to keep (0=unlimited)")
    p_filter.add_argument("--keywords", nargs="*", help="For papers without score, require these keywords in title")

    # export
    p_export = sub.add_parser("export", help="Export database")
    p_export.add_argument("--input", required=True, help="Paper DB JSONL file")
    p_export.add_argument("--format", choices=["csv", "jsonl"], default="csv")
    p_export.add_argument("--output", "-o", help="Output file (default: stdout)")

    args = parser.parse_args()

    if args.command == "merge":
        merge_databases(args.inputs, args.output, args.threshold)

    elif args.command == "search":
        results = search_db(args.input, args.query, args.field)
        for rec in results:
            print(json.dumps(rec, ensure_ascii=False))
        print(f"Found {len(results)} matches", file=sys.stderr)

    elif args.command == "tag":
        tag_papers(args.input, args.ids, args.tags)

    elif args.command == "add":
        record = json.loads(args.record)
        records = load_jsonl(args.input)
        records.append(record)
        records = deduplicate(records)
        save_jsonl(records, args.input)
        print(f"Added record, DB now has {len(records)} papers", file=sys.stderr)

    elif args.command == "filter":
        filter_db(args.input, args.output, min_score=args.min_score,
                  max_papers=args.max_papers, require_keywords=args.keywords)

    elif args.command == "stats":
        stats = compute_stats(args.input)
        print(json.dumps(stats, indent=2))

    elif args.command == "export":
        if args.format == "csv":
            output = export_csv(args.input)
        else:
            records = load_jsonl(args.input)
            output = "\n".join(json.dumps(r, ensure_ascii=False) for r in records)

        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                f.write(output)
            print(f"Exported to {args.output}", file=sys.stderr)
        else:
            print(output)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Search arxiv via the Atom API and output JSONL paper metadata.

Self-contained: uses only stdlib (urllib, xml.etree).

Usage:
    python search_arxiv.py --query "long context reasoning" --max-results 50
    python search_arxiv.py --query "protein language model" --categories q-bio.BM cs.LG --max-results 100
    python search_arxiv.py --query "LLM agent" --sort-by lastUpdatedDate --start-date 2024-01-01
"""

import argparse
import json
import sys
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from datetime import datetime

ARXIV_API = "http://export.arxiv.org/api/query"
NS = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}

SORT_MAP = {
    "relevance": "relevance",
    "lastUpdatedDate": "lastUpdatedDate",
    "submittedDate": "submittedDate",
}


def build_query(keywords: str, categories: list[str] | None = None) -> str:
    """Build an arxiv search query string."""
    parts = []
    # keyword search across all fields
    parts.append(f"all:{keywords}")
    if categories:
        cat_query = " OR ".join(f"cat:{c}" for c in categories)
        parts.append(f"({cat_query})")
    return " AND ".join(parts)


def fetch_results(
    query: str,
    start: int,
    max_results: int,
    sort_by: str = "relevance",
    sort_order: str = "descending",
) -> bytes:
    """Fetch a page of results from the arxiv API."""
    params = {
        "search_query": query,
        "start": start,
        "max_results": max_results,
        "sortBy": sort_by,
        "sortOrder": sort_order,
    }
    url = f"{ARXIV_API}?{urllib.parse.urlencode(params)}"
    req = urllib.request.Request(url, headers={"User-Agent": "deep-research/1.0"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        return resp.read()


def parse_entry(entry: ET.Element) -> dict:
    """Parse a single Atom entry into a paper record."""
    def text(tag: str, ns: str = "atom") -> str:
        el = entry.find(f"{ns}:{tag}", NS) if ns else entry.find(tag)
        return el.text.strip() if el is not None and el.text else ""

    # Extract arxiv ID from the entry id URL
    entry_id = text("id")
    arxiv_id = entry_id.split("/abs/")[-1] if "/abs/" in entry_id else entry_id

    # Authors
    authors = []
    for author_el in entry.findall("atom:author", NS):
        name_el = author_el.find("atom:name", NS)
        if name_el is not None and name_el.text:
            authors.append(name_el.text.strip())

    # Categories
    categories = []
    for cat_el in entry.findall("arxiv:primary_category", NS):
        term = cat_el.get("term", "")
        if term:
            categories.append(term)
    for cat_el in entry.findall("atom:category", NS):
        term = cat_el.get("term", "")
        if term and term not in categories:
            categories.append(term)

    # PDF link
    pdf_url = ""
    for link_el in entry.findall("atom:link", NS):
        if link_el.get("title") == "pdf":
            pdf_url = link_el.get("href", "")
            break
    if not pdf_url and arxiv_id:
        pdf_url = f"https://arxiv.org/pdf/{arxiv_id}"

    # Comment (often contains page count, conference info)
    comment = text("comment", "arxiv")

    # Abstract: normalize whitespace
    abstract = " ".join(text("summary").split())

    published = text("published")
    year = int(published[:4]) if len(published) >= 4 else None

    return {
        "arxiv_id": arxiv_id,
        "title": " ".join(text("title").split()),
        "authors": authors,
        "abstract": abstract,
        "year": year,
        "published": published,
        "updated": text("updated"),
        "categories": categories,
        "pdf_url": pdf_url,
        "comment": comment,
        "source": "arxiv",
    }


def search(
    keywords: str,
    categories: list[str] | None = None,
    max_results: int = 50,
    sort_by: str = "relevance",
    start_date: str | None = None,
    end_date: str | None = None,
) -> list[dict]:
    """Run a full paginated search and return deduplicated results."""
    query = build_query(keywords, categories)
    page_size = min(max_results, 100)  # arxiv max per request
    all_papers = []
    seen_ids = set()

    for start in range(0, max_results, page_size):
        fetch_count = min(page_size, max_results - start)
        try:
            xml_data = fetch_results(query, start, fetch_count, sort_by)
        except Exception as e:
            print(f"Warning: fetch failed at offset {start}: {e}", file=sys.stderr)
            break

        root = ET.fromstring(xml_data)
        entries = root.findall("atom:entry", NS)
        if not entries:
            break

        for entry in entries:
            paper = parse_entry(entry)
            if not paper["title"] or paper["arxiv_id"] in seen_ids:
                continue

            # Date filtering
            if start_date or end_date:
                pub = paper["published"][:10]  # YYYY-MM-DD
                if start_date and pub < start_date:
                    continue
                if end_date and pub > end_date:
                    continue

            seen_ids.add(paper["arxiv_id"])
            all_papers.append(paper)

        # Respect rate limit: 1 request per 3 seconds
        if start + page_size < max_results:
            time.sleep(3)

    return all_papers


def main():
    parser = argparse.ArgumentParser(description="Search arxiv and output JSONL")
    parser.add_argument("--query", required=True, help="Search keywords")
    parser.add_argument("--max-results", type=int, default=50, help="Max papers to return")
    parser.add_argument("--categories", nargs="*", help="arxiv categories (e.g. cs.AI cs.CL q-bio.BM)")
    parser.add_argument("--sort-by", choices=list(SORT_MAP.keys()), default="relevance")
    parser.add_argument("--start-date", help="Filter: earliest publication date (YYYY-MM-DD)")
    parser.add_argument("--end-date", help="Filter: latest publication date (YYYY-MM-DD)")
    parser.add_argument("--output", "-o", help="Output file (default: stdout)")
    args = parser.parse_args()

    papers = search(
        keywords=args.query,
        categories=args.categories,
        max_results=args.max_results,
        sort_by=SORT_MAP[args.sort_by],
        start_date=args.start_date,
        end_date=args.end_date,
    )

    out = open(args.output, "w") if args.output else sys.stdout
    try:
        for paper in papers:
            out.write(json.dumps(paper, ensure_ascii=False) + "\n")
    finally:
        if args.output:
            out.close()

    print(f"Found {len(papers)} papers", file=sys.stderr)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""Search Semantic Scholar Graph API and output JSONL paper metadata.

Self-contained: uses only stdlib (urllib, json).

Usage:
    python search_semantic_scholar.py --query "long horizon reasoning" --max-results 100
    python search_semantic_scholar.py --query "protein language model" --min-citations 10 --year-range 2020-2026
    python search_semantic_scholar.py --query "LLM agent planning" --venue NeurIPS ICML --max-results 50
"""

import argparse
import json
import sys
import time
import urllib.parse
import urllib.request

S2_API = "https://api.semanticscholar.org/graph/v1"
FIELDS = "title,authors,abstract,year,venue,citationCount,externalIds,url,referenceCount,publicationDate"
SEARCH_LIMIT = 100  # S2 max per request

# Top-tier AI/ML conferences (peer-reviewed)
TOP_CONFERENCES = {
    # ML core
    "NeurIPS", "ICML", "ICLR",
    # NLP
    "ACL", "EMNLP", "NAACL", "EACL", "COLING",
    # AI general
    "AAAI", "IJCAI",
    # Vision (occasionally relevant)
    "CVPR", "ICCV", "ECCV",
    # IR / data mining
    "KDD", "WWW", "SIGIR",
    # Robotics / agents
    "ICRA", "CoRL",
}

# Normalized aliases for S2 venue matching
VENUE_ALIASES = {
    "neurips": "NeurIPS", "nips": "NeurIPS",
    "icml": "ICML",
    "iclr": "ICLR",
    "acl": "ACL",
    "emnlp": "EMNLP",
    "naacl": "NAACL",
    "eacl": "EACL",
    "coling": "COLING",
    "aaai": "AAAI",
    "ijcai": "IJCAI",
    "cvpr": "CVPR",
    "iccv": "ICCV",
    "eccv": "ECCV",
    "kdd": "KDD",
    "www": "WWW",
    "sigir": "SIGIR",
    "icra": "ICRA",
    "corl": "CoRL",
}


def is_peer_reviewed(venue: str) -> bool:
    """Check if a paper's venue is a recognized peer-reviewed conference."""
    if not venue:
        return False
    venue_lower = venue.lower()
    for alias in VENUE_ALIASES:
        if alias in venue_lower:
            return True
    # Also check for "journal" or "transactions" as peer-reviewed
    if any(kw in venue_lower for kw in ("journal", "transactions", "review")):
        return True
    return False


def normalize_venue(venue: str) -> str:
    """Normalize venue name to canonical form."""
    if not venue:
        return ""
    venue_lower = venue.lower()
    for alias, canonical in VENUE_ALIASES.items():
        if alias in venue_lower:
            return canonical
    return venue


def s2_request(url: str, api_key: str | None = None) -> dict:
    """Make a request to the Semantic Scholar API with retry logic."""
    headers = {"User-Agent": "deep-research/1.0"}
    if api_key:
        headers["x-api-key"] = api_key

    req = urllib.request.Request(url, headers=headers)

    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                return json.loads(resp.read())
        except urllib.error.HTTPError as e:
            if e.code == 429:
                wait = 2 ** (attempt + 1)
                print(f"Rate limited, waiting {wait}s...", file=sys.stderr)
                time.sleep(wait)
                continue
            raise
        except Exception:
            if attempt < 2:
                time.sleep(1)
                continue
            raise
    return {}


def parse_paper(data: dict) -> dict | None:
    """Parse an S2 paper response into our standard record format."""
    if not data or not data.get("title"):
        return None

    authors = []
    for a in data.get("authors", []) or []:
        name = a.get("name", "")
        if name:
            authors.append(name)

    external_ids = data.get("externalIds", {}) or {}
    arxiv_id = external_ids.get("ArXiv", "")

    pdf_url = ""
    if arxiv_id:
        pdf_url = f"https://arxiv.org/pdf/{arxiv_id}"

    abstract = data.get("abstract", "") or ""

    venue = data.get("venue", "") or ""
    reviewed = is_peer_reviewed(venue)

    return {
        "paperId": data.get("paperId", ""),
        "arxiv_id": arxiv_id,
        "title": data["title"],
        "authors": authors,
        "abstract": " ".join(abstract.split()),
        "year": data.get("year"),
        "venue": venue,
        "venue_normalized": normalize_venue(venue),
        "peer_reviewed": reviewed,
        "citationCount": data.get("citationCount", 0) or 0,
        "referenceCount": data.get("referenceCount", 0) or 0,
        "url": data.get("url", ""),
        "publicationDate": data.get("publicationDate", ""),
        "pdf_url": pdf_url,
        "source": "semantic_scholar",
    }


def search_papers(
    query: str,
    max_results: int = 100,
    year_range: str | None = None,
    min_citations: int = 0,
    venue_filter: list[str] | None = None,
    peer_reviewed_only: bool = False,
    api_key: str | None = None,
) -> list[dict]:
    """Search for papers and return deduplicated results.

    If peer_reviewed_only=True, only returns papers from recognized conferences/journals.
    This fetches more results from S2 to compensate for filtering.
    """
    all_papers = []
    seen_ids = set()

    # When filtering peer-reviewed, fetch more to compensate
    fetch_multiplier = 3 if peer_reviewed_only else 1
    fetch_max = max_results * fetch_multiplier

    offset = 0
    while offset < fetch_max and len(all_papers) < max_results:
        limit = min(SEARCH_LIMIT, fetch_max - offset)
        params = {
            "query": query,
            "offset": offset,
            "limit": limit,
            "fields": FIELDS,
        }
        if year_range:
            params["year"] = year_range

        url = f"{S2_API}/paper/search?{urllib.parse.urlencode(params)}"

        try:
            resp = s2_request(url, api_key)
        except Exception as e:
            print(f"Warning: search failed at offset {offset}: {e}", file=sys.stderr)
            break

        papers = resp.get("data", [])
        if not papers:
            break

        for item in papers:
            if len(all_papers) >= max_results:
                break

            paper = parse_paper(item)
            if not paper:
                continue

            # Dedup by paperId
            pid = paper["paperId"]
            if pid in seen_ids:
                continue
            seen_ids.add(pid)

            # Citation filter
            if paper["citationCount"] < min_citations:
                continue

            # Venue filter (explicit list)
            if venue_filter:
                paper_venue = (paper.get("venue", "") or "").lower()
                if not any(v.lower() in paper_venue for v in venue_filter):
                    continue

            # Peer-reviewed filter
            if peer_reviewed_only and not paper.get("peer_reviewed", False):
                continue

            all_papers.append(paper)

        total = resp.get("total", 0)
        offset += limit
        if offset >= total:
            break

        # Rate limit: ~10 requests per second for public API
        time.sleep(0.5)

    return all_papers


def get_paper_details(paper_id: str, api_key: str | None = None) -> dict | None:
    """Get detailed info for a single paper by S2 paperId or arxiv:<id>."""
    url = f"{S2_API}/paper/{urllib.parse.quote(paper_id, safe='')}?fields={FIELDS}"
    try:
        data = s2_request(url, api_key)
        return parse_paper(data)
    except Exception as e:
        print(f"Warning: failed to fetch {paper_id}: {e}", file=sys.stderr)
        return None


def get_citations(paper_id: str, max_results: int = 50, api_key: str | None = None) -> list[dict]:
    """Get papers that cite the given paper."""
    url = f"{S2_API}/paper/{urllib.parse.quote(paper_id, safe='')}/citations?fields={FIELDS}&limit={min(max_results, 1000)}"
    try:
        resp = s2_request(url, api_key)
        results = []
        for item in resp.get("data", []):
            citing = item.get("citingPaper", {})
            paper = parse_paper(citing)
            if paper:
                results.append(paper)
        return results[:max_results]
    except Exception as e:
        print(f"Warning: failed to fetch citations for {paper_id}: {e}", file=sys.stderr)
        return []


def get_references(paper_id: str, max_results: int = 50, api_key: str | None = None) -> list[dict]:
    """Get papers referenced by the given paper."""
    url = f"{S2_API}/paper/{urllib.parse.quote(paper_id, safe='')}/references?fields={FIELDS}&limit={min(max_results, 1000)}"
    try:
        resp = s2_request(url, api_key)
        results = []
        for item in resp.get("data", []):
            cited = item.get("citedPaper", {})
            paper = parse_paper(cited)
            if paper:
                results.append(paper)
        return results[:max_results]
    except Exception as e:
        print(f"Warning: failed to fetch references for {paper_id}: {e}", file=sys.stderr)
        return []


def main():
    # Handle --list-conferences early (no other args needed)
    if "--list-conferences" in sys.argv:
        print("Recognized top conferences:")
        for conf in sorted(TOP_CONFERENCES):
            print(f"  {conf}")
        sys.exit(0)

    parser = argparse.ArgumentParser(description="Search Semantic Scholar and output JSONL")
    parser.add_argument("--query", required=True, help="Search keywords")
    parser.add_argument("--max-results", type=int, default=100, help="Max papers to return")
    parser.add_argument("--min-citations", type=int, default=0, help="Minimum citation count")
    parser.add_argument("--year-range", help="Year range filter (e.g. 2020-2026)")
    parser.add_argument("--venue", nargs="*", help="Venue filter (e.g. NeurIPS ICML)")
    parser.add_argument("--peer-reviewed-only", action="store_true",
                        help="Only return papers from peer-reviewed conferences/journals")
    parser.add_argument("--top-conferences", action="store_true",
                        help="Shorthand for --venue with all top AI conferences")
    parser.add_argument("--list-conferences", action="store_true",
                        help="Print the list of recognized top conferences and exit")
    parser.add_argument("--api-key", help="S2 API key (optional, increases rate limits)")
    parser.add_argument("--citations-of", help="Get papers citing this paper ID")
    parser.add_argument("--references-of", help="Get papers referenced by this paper ID")
    parser.add_argument("--output", "-o", help="Output file (default: stdout)")
    args = parser.parse_args()

    if args.list_conferences:
        print("Recognized top conferences:")
        for conf in sorted(TOP_CONFERENCES):
            print(f"  {conf}")
        sys.exit(0)

    # --top-conferences expands to venue filter with all top conferences
    venue_filter = args.venue
    if args.top_conferences:
        venue_filter = list(TOP_CONFERENCES)

    if args.citations_of:
        papers = get_citations(args.citations_of, args.max_results, args.api_key)
    elif args.references_of:
        papers = get_references(args.references_of, args.max_results, args.api_key)
    else:
        papers = search_papers(
            query=args.query,
            max_results=args.max_results,
            year_range=args.year_range,
            min_citations=args.min_citations,
            venue_filter=venue_filter,
            peer_reviewed_only=args.peer_reviewed_only,
            api_key=args.api_key,
        )

    out = open(args.output, "w") if args.output else sys.stdout
    try:
        for paper in papers:
            out.write(json.dumps(paper, ensure_ascii=False) + "\n")
    finally:
        if args.output:
            out.close()

    print(f"Found {len(papers)} papers", file=sys.stderr)


if __name__ == "__main__":
    main()