stack/dev/scripts/snowball_citations.py

"""Forward/backward snowball citation chasing for skin substitutes.

Uses NCBI elink to find articles that cite (forward snowball) the
top-cited RCTs and meta-analyses in our PubMed collection. Identifies
PMIDs not already in bib.sqlite and fetches their metadata.

Addresses #237 checklist: "Forward/backward snowball from included
studies — identify missed references".

Usage:
    uv run python dev/scripts/snowball_citations.py
    uv run python dev/scripts/snowball_citations.py --dry-run
"""

from __future__ import annotations

import argparse
import os
import time
import xml.etree.ElementTree as ET
from datetime import datetime

import httpx

from bib.item import Source
from bib.store import Store

EUTILS_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
API_KEY = os.environ.get("NCBI_API_KEY", "")
TOOL_NAME = "stack-skin-subs-snowball"
TOOL_EMAIL = "dev@localhost"
RATE_LIMIT = 0.34 if not API_KEY else 0.1


def _params(**kw: str) -> dict[str, str]:
    base = {"tool": TOOL_NAME, "email": TOOL_EMAIL}
    if API_KEY:
        base["api_key"] = API_KEY
    base.update(kw)
    return base


# ---------------------------------------------------------------------------
# elink: find citing articles (forward snowball)
# ---------------------------------------------------------------------------


def elink_cited_by(pmids: list[str], batch_size: int = 50) -> dict[str, list[str]]:
    """For each PMID, find PMIDs that cite it (forward snowball).

    Uses elink with linkname=pubmed_pubmed_citedin.
    Returns {source_pmid: [citing_pmid, ...]}.
    """
    result: dict[str, list[str]] = {}
    for i in range(0, len(pmids), batch_size):
        batch = pmids[i : i + batch_size]
        # `id` must repeat once per PMID: a single comma-joined value
        # makes elink merge every seed into ONE LinkSet, so all citing
        # articles get attributed to the first seed (bit the
        # palliative-rfi snowball tagging on 2026-08-18).
        params = list(
            _params(
                dbfrom="pubmed",
                db="pubmed",
                linkname="pubmed_pubmed_citedin",
                retmode="xml",
            ).items()
        ) + [("id", p) for p in batch]
        time.sleep(RATE_LIMIT)
        try:
            resp = httpx.get(f"{EUTILS_BASE}/elink.fcgi", params=params, timeout=60)
            resp.raise_for_status()
            root = ET.fromstring(resp.text)  # noqa: S314
            for linkset in root.findall(".//LinkSet"):
                id_el = linkset.find("IdList/Id")
                if id_el is None or not id_el.text:
                    continue
                src_pmid = id_el.text
                citing = []
                for link_db in linkset.findall(".//LinkSetDb"):
                    ln = link_db.find("LinkName")
                    if ln is not None and ln.text == "pubmed_pubmed_citedin":
                        for lid in link_db.findall("Link/Id"):
                            if lid.text:
                                citing.append(lid.text)
                result[src_pmid] = citing
        except Exception as exc:
            print(f"  elink error batch {i // batch_size + 1}: {exc}")

        if (i // batch_size + 1) % 10 == 0:
            print(f"  elink: processed {i + len(batch)}/{len(pmids)} seeds")

    return result


# ---------------------------------------------------------------------------
# efetch: get article metadata for new PMIDs
# ---------------------------------------------------------------------------


def _text(el: ET.Element | None, path: str, default: str = "") -> str:
    """Get all text including mixed-content children (sup, i, sub, etc.)."""
    if el is None:
        return default
    node = el.find(path)
    if node is None:
        return default
    return "".join(node.itertext()).strip() or default


def efetch_basic(pmids: list[str], batch_size: int = 100) -> list[dict]:
    """Fetch basic article info for a list of PMIDs."""
    articles: list[dict] = []
    for i in range(0, len(pmids), batch_size):
        batch = pmids[i : i + batch_size]
        params = _params(
            db="pubmed",
            id=",".join(batch),
            rettype="xml",
            retmode="xml",
        )
        for attempt in range(3):
            time.sleep(RATE_LIMIT * (attempt + 1))
            try:
                resp = httpx.get(
                    f"{EUTILS_BASE}/efetch.fcgi",
                    params=params,
                    timeout=120,
                )
                resp.raise_for_status()
                root = ET.fromstring(resp.text)  # noqa: S314
                for art_el in root.findall(".//PubmedArticle"):
                    citation = art_el.find(".//MedlineCitation")
                    if citation is None:
                        continue
                    pmid = _text(citation, "PMID")
                    article_el = citation.find("Article")
                    if article_el is None:
                        continue
                    title = _text(article_el, "ArticleTitle")

                    # Abstract
                    abstract_parts = []
                    abstract_el = article_el.find("Abstract")
                    if abstract_el is not None:
                        for at in abstract_el.findall("AbstractText"):
                            label = at.get("Label", "")
                            text = "".join(at.itertext()).strip()
                            if label:
                                abstract_parts.append(f"{label}: {text}")
                            else:
                                abstract_parts.append(text)

                    # Authors
                    authors = []
                    author_list = article_el.find("AuthorList")
                    if author_list is not None:
                        for au in author_list.findall("Author"):
                            last = _text(au, "LastName")
                            fore = _text(au, "ForeName")
                            if last:
                                authors.append(f"{last} {fore}".strip())

                    # Journal + year
                    journal_el = article_el.find("Journal")
                    journal = (
                        _text(journal_el, "Title") if journal_el is not None else ""
                    )
                    year = ""
                    pub_date = article_el.find(".//PubDate")
                    if pub_date is not None:
                        year = _text(pub_date, "Year")
                        if not year:
                            md = _text(pub_date, "MedlineDate")
                            if md:
                                year = md[:4]

                    # DOI
                    doi = ""
                    for id_el in art_el.findall(".//ArticleId"):
                        if id_el.get("IdType") == "doi":
                            doi = id_el.text or ""
                            break

                    # Pub types
                    pub_types = []
                    for pt in article_el.findall(".//PublicationType"):
                        if pt.text:
                            pub_types.append(pt.text)

                    articles.append(
                        {
                            "pmid": pmid,
                            "title": title,
                            "abstract": "\n\n".join(abstract_parts),
                            "authors": authors,
                            "journal": journal,
                            "year": year,
                            "doi": doi,
                            "pub_types": pub_types,
                        }
                    )
                break
            except (httpx.RemoteProtocolError, httpx.ReadTimeout) as exc:
                if attempt < 2:
                    time.sleep(2 ** (attempt + 1))
                else:
                    print(f"  efetch skip batch {i // batch_size + 1}: {exc}")

        if (i // batch_size + 1) % 10 == 0:
            print(f"  efetch: {len(articles)} articles so far")

    return articles


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description="Snowball citation chasing")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--seed-limit",
        type=int,
        default=200,
        help="Max seed articles for forward snowball",
    )
    args = parser.parse_args()

    print("=" * 70)
    print("Snowball Citation Chasing: Skin Substitutes")
    print(f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}")
    print("=" * 70)

    store = Store()
    con = store._con()

    # --- Select seed articles: top RCTs and meta-analyses ---
    print("\n--- Selecting seed articles ---")
    # Get PMIDs for RCTs and meta-analyses
    seed_rows = con.execute(
        """SELECT DISTINCT i.extra, i.url
           FROM items i
           JOIN item_tags it1 ON i.id = it1.item_id
           JOIN tags t1 ON it1.tag_id = t1.id
           JOIN item_tags it2 ON i.id = it2.item_id
           JOIN tags t2 ON it2.tag_id = t2.id
           WHERE t1.name = 'module:skin-subs'
             AND t2.name IN ('type:rct', 'type:meta-analysis', 'type:review')
             AND i.url LIKE '%pubmed%'"""
    ).fetchall()

    seed_pmids = []
    for row in seed_rows:
        extra = row["extra"] or ""
        for line in extra.split("\n"):
            if line.startswith("PMID:"):
                pmid = line[5:].strip()
                if pmid:
                    seed_pmids.append(pmid)
                break

    # Limit seeds
    seed_pmids = seed_pmids[: args.seed_limit]
    print(f"  Seed articles (RCTs + meta-analyses + reviews): {len(seed_pmids)}")

    # --- Get existing PMIDs to deduplicate ---
    all_existing = con.execute(
        """SELECT DISTINCT i.extra FROM items i
           JOIN item_tags it ON i.id = it.item_id
           JOIN tags t ON it.tag_id = t.id
           WHERE t.name = 'module:skin-subs'
             AND i.url LIKE '%pubmed%'"""
    ).fetchall()

    existing_pmids = set()
    for row in all_existing:
        extra = row["extra"] or ""
        for line in extra.split("\n"):
            if line.startswith("PMID:"):
                existing_pmids.add(line[5:].strip())
                break
    print(f"  Existing PMIDs in bib.sqlite: {len(existing_pmids)}")

    # --- Forward snowball ---
    print("\n--- Forward snowball (cited-by) ---")
    cited_by = elink_cited_by(seed_pmids)

    all_citing = set()
    for src, citing_list in cited_by.items():
        all_citing.update(citing_list)

    new_pmids = all_citing - existing_pmids
    print(f"  Total citing articles found: {len(all_citing)}")
    print(f"  Already in collection: {len(all_citing - new_pmids)}")
    print(f"  New articles to add: {len(new_pmids)}")

    if not new_pmids:
        print("\nNo new articles found. Done.")
        store.close()
        return

    if args.dry_run:
        print(f"\n[DRY RUN] Would fetch and store {len(new_pmids)} new articles")
        store.close()
        return

    # --- Fetch metadata for new PMIDs ---
    print(f"\n--- Fetching metadata for {len(new_pmids)} new articles ---")
    new_articles = efetch_basic(sorted(new_pmids))
    print(f"  Fetched {len(new_articles)} articles")

    # --- Relevance filter: must mention skin/wound in title or abstract ---
    skin_keywords = [
        "skin substitute",
        "skin substitutes",
        "wound",
        "ulcer",
        "biological dressing",
        "tissue product",
        "graft",
        "dermal",
        "epidermal",
        "bioengineered",
    ]
    relevant = []
    for art in new_articles:
        text = f"{art['title']} {art['abstract']}".lower()
        if any(kw in text for kw in skin_keywords):
            relevant.append(art)

    print(f"  Relevant (mention skin/wound): {len(relevant)}")
    print(f"  Filtered out: {len(new_articles) - len(relevant)}")

    # --- Store in bib.sqlite ---
    print(f"\n--- Storing {len(relevant)} snowball articles ---")
    created = 0
    for art in relevant:
        author_str = "; ".join(art["authors"][:10])
        extra_parts = [f"PMID: {art['pmid']}"]
        if art["doi"]:
            extra_parts.append(f"DOI: {art['doi']}")
        if author_str:
            extra_parts.append(f"Authors: {author_str}")
        if art["journal"]:
            extra_parts.append(f"Journal: {art['journal']}")
        if art["pub_types"]:
            extra_parts.append(f"PubTypes: {'; '.join(art['pub_types'])}")

        tags = ["module:skin-subs", "source:pubmed", "source:snowball"]
        if art["year"]:
            tags.append(f"year:{art['year']}")

        # Classify type
        pt_lower = [p.lower() for p in art["pub_types"]]
        if "meta-analysis" in pt_lower:
            tags.append("type:meta-analysis")
        elif "randomized controlled trial" in pt_lower:
            tags.append("type:rct")
        elif "review" in pt_lower or "systematic review" in pt_lower:
            tags.append("type:review")

        source = Source(
            title=art["title"],
            url=f"https://pubmed.ncbi.nlm.nih.gov/{art['pmid']}/",
            date_published=art["year"],
            institution=art["journal"],
            abstract=art["abstract"],
            doc_type="journal-article",
            tags=tags,
            extra="\n".join(extra_parts),
        )
        store.upsert(source)
        created += 1

    print(f"  Created/updated: {created}")

    # Final count
    total = con.execute(
        """SELECT count(DISTINCT i.id) FROM items i
           JOIN item_tags it ON i.id = it.item_id
           JOIN tags t ON it.tag_id = t.id
           WHERE t.name = 'module:skin-subs'"""
    ).fetchone()[0]
    snowball_count = con.execute(
        """SELECT count(DISTINCT i.id) FROM items i
           JOIN item_tags it ON i.id = it.item_id
           JOIN tags t ON it.tag_id = t.id
           WHERE t.name = 'source:snowball'"""
    ).fetchone()[0]
    print(f"\n  Total module:skin-subs items: {total}")
    print(f"  Snowball additions: {snowball_count}")

    store.close()
    print("\nDone.")


if __name__ == "__main__":
    main()