backend/astroparse_api/references.py

"""Parse the raw references section text into structured entries with corpus matching."""
import re
from dataclasses import dataclass

_YEAR_RE = re.compile(r"\b((?:19|20)\d{2})\b")
_ARXIV_RE = re.compile(r"arXiv[:\s](\d{4}\.\d{4,5})", re.I)
_JOURNAL_RE = re.compile(
    r"\b(ApJL|ApJS|ApJ|AJ|MNRAS|A&A|ARA&A|ARAA|PASP|Nature|Science|JCAP)\b"
)
_VOL_PAGE_RE = re.compile(r",\s*(\d+),\s*(\w+)")

# ADS journal abbreviations (5 chars, dot-padded on right)
_JOURNAL_ADS = {
    "ApJ":    "ApJ..",
    "ApJL":   "ApJ..",   # same journal code; page prefix 'L' distinguishes
    "ApJS":   "ApJS.",
    "AJ":     "AJ...",
    "MNRAS":  "MNRAS",
    "A&A":    "A&A..",
    "ARA&A":  "ARA&A",
    "ARAA":   "ARA&A",
    "PASP":   "PASP.",
    "Nature": "Natur",
    "Science": "Sci..",
    "JCAP":   "JCAP.",
}

# Split on lines that begin a new reference entry (capitalized surname)
_ENTRY_SPLIT = re.compile(r"\n(?=[A-Z][a-zA-ZÀ-ÿ\-']{1,}(?:[,\s&]|$))", re.M)


@dataclass
class ParsedRef:
    id: str
    raw: str
    short: str
    year: int | None
    bibcode: str | None
    arxiv: str | None
    corpus_match: bool


def _vol_pad(vol: str) -> str:
    """Left dot-pad volume to 4 chars (ADS convention).

    ADS pads volumes on the left with dots: '873' → '.873'.
    """
    return vol[:4].rjust(4, ".")


def _page_pad(page: str) -> str:
    """Left dot-pad page number to 4 chars (ADS convention).

    ADS pads page numbers on the left with dots: '14' → '..14'.
    """
    return page[:4].rjust(4, ".")


def _build_bibcode(year: str, journal: str, vol: str, page: str, initial: str) -> str | None:
    """Construct an ADS bibcode from parsed fields. Returns None for unknown journals.

    ADS bibcode format (19 chars): YYYY JJJJJ VVVV M PPPP A
      - YYYY: 4-char year
      - JJJJJ: 5-char journal abbreviation (right dot-padded)
      - VVVV: 4-char volume (left dot-padded with '.')
      - M: 1-char page-type flag ('L' for Letters, '.' for normal articles)
      - PPPP: 4-char page number (left dot-padded with '.')
      - A: 1-char first author initial (uppercase)

    Examples:
      2018ApJ...873L..14C  → year=2018, j=ApJ.., vol=.873, M=L, page=..14, init=C
      2018MNRAS.489.3023S  → year=2018, j=MNRAS, vol=.489, M=., page=3023, init=S
    """
    abbrev = _JOURNAL_ADS.get(journal)
    if not abbrev:
        return None

    vol_part = _vol_pad(vol)

    # Handle letter pages (e.g. "L14" for ApJL)
    if journal == "ApJL" and page.startswith("L"):
        page_num = page[1:]  # strip the 'L' prefix
        page_part = "L" + _page_pad(page_num)
    else:
        page_part = "." + _page_pad(page)

    return f"{year}{abbrev}{vol_part}{page_part}{initial.upper()}"


def _extract_surname(text: str) -> str:
    m = re.match(r"([A-Z][a-zA-ZÀ-ÿ\-']+)", text)
    return m.group(1) if m else ""


def parse_references(
    raw_refs: str,
    arxiv_ids: set[str],
    bibcodes: set[str],
) -> tuple[list[ParsedRef], dict[str, list[str]]]:
    """Parse raw reference section text into ParsedRef list + citeIndex.

    citeIndex maps '{surname_lower}:{year}' → [ref_id, ...].
    Corpus matching: arXiv id exact match first, then constructed bibcode.
    Misses degrade to display-only (corpus_match=False), never wrong matches.
    """
    if not raw_refs or not raw_refs.strip():
        return [], {}

    entries = _ENTRY_SPLIT.split(raw_refs.strip())
    refs: list[ParsedRef] = []
    cite_index: dict[str, list[str]] = {}

    for i, raw in enumerate(entries):
        raw = raw.strip()
        if not raw or len(raw) < 15:
            continue

        ref_id = f"r{i + 1}"

        ym = _YEAR_RE.search(raw)
        year = int(ym.group(1)) if ym else None

        am = _ARXIV_RE.search(raw)
        arxiv = am.group(1) if am else None

        surname = _extract_surname(raw)
        short = f"{surname} {year}" if surname and year else surname or str(year or "")

        # Corpus matching
        corpus_match = False
        bibcode = None

        if arxiv and arxiv in arxiv_ids:
            corpus_match = True

        if not corpus_match and year:
            jm = _JOURNAL_RE.search(raw)
            vm = _VOL_PAGE_RE.search(raw)
            if jm and vm:
                bc = _build_bibcode(
                    str(year), jm.group(1), vm.group(1), vm.group(2),
                    surname[0] if surname else "X",
                )
                if bc and bc in bibcodes:
                    corpus_match = True
                    bibcode = bc

        ref = ParsedRef(
            id=ref_id, raw=raw, short=short, year=year,
            bibcode=bibcode, arxiv=arxiv, corpus_match=corpus_match,
        )
        refs.append(ref)

        if surname and year:
            key = f"{surname.lower()}:{year}"
            cite_index.setdefault(key, []).append(ref_id)

    return refs, cite_index