"""Renumber the manuscript's references to first-appearance order (PNAS style).

Reads paper/pnas/main.md, finds every parenthesised citation group in the text above
"## References", derives the order in which references first appear, and rewrites the citation
groups in main.md, si.md, and the figure captions in build.py, then reorders the reference list.
Citation groups are parentheses containing only reference numbers, commas, en-dash ranges, an optional
"cf. " prefix, or a prose prefix ending in a semicolon ("...; 11, 12"). Four-digit numbers (years)
never match, and any number above the list length is reported and left alone.

Usage:  python paper/pnas/renumber_refs.py            # dry run: mapping + per-file counts
        python paper/pnas/renumber_refs.py --apply    # rewrite the three files in place
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
MAIN, SI, BUILD = (ROOT / "paper" / "pnas" / n for n in ("main.md", "si.md", "build.py"))
REF_HEADER = "## References"
CIT = re.compile(
    r"\((?P<pre>[^()]*?;\s*)?(?P<cf>cf\.\s*)?"
    r"(?P<nums>\d{1,3}(?:\s*[–-]\s*\d{1,3})?(?:,\s*\d{1,3}(?:\s*[–-]\s*\d{1,3})?)*)\)"
)
REF_LINE = re.compile(r"^(\d+)\. (.*)$")


def expand(nums: str) -> list[int]:
    out: list[int] = []
    for part in re.split(r",\s*", nums):
        if re.search(r"[–-]", part):
            a, b = (int(x) for x in re.split(r"\s*[–-]\s*", part))
            out.extend(range(a, b + 1))
        else:
            out.append(int(part))
    return out


def compress(nums: list[int]) -> str:
    """Ascending, with runs of three or more collapsed to an en-dash range."""
    nums = sorted(set(nums))
    runs: list[list[int]] = []
    for n in nums:
        if runs and n == runs[-1][-1] + 1:
            runs[-1].append(n)
        else:
            runs.append([n])
    return ", ".join(f"{r[0]}–{r[-1]}" if len(r) >= 3 else ", ".join(map(str, r)) for r in runs)


def split_main(text: str) -> tuple[str, list[tuple[int, str]]]:
    body, _, refs = text.partition(REF_HEADER)
    entries = [(int(m.group(1)), m.group(2)) for line in refs.splitlines() if (m := REF_LINE.match(line))]
    return body, entries


def first_appearance(body: str, n_refs: int) -> list[int]:
    order: list[int] = []
    for m in CIT.finditer(body):
        for n in expand(m.group("nums")):
            if n <= n_refs and n not in order:
                order.append(n)
    return order


def rewrite(text: str, mapping: dict[int, int], n_refs: int, label: str) -> tuple[str, int, list[str]]:
    count, suspicious = 0, []

    def sub(m: re.Match) -> str:
        nonlocal count
        nums = expand(m.group("nums"))
        if any(n > n_refs or n < 1 for n in nums):
            suspicious.append(m.group(0))
            return m.group(0)
        count += 1
        return f"({m.group('pre') or ''}{m.group('cf') or ''}{compress([mapping[n] for n in nums])})"

    return CIT.sub(sub, text), count, suspicious


def main(apply: bool) -> int:
    main_text = MAIN.read_text()
    body, entries = split_main(main_text)
    n_refs = len(entries)
    assert [n for n, _ in entries] == list(range(1, n_refs + 1)), "reference list is not 1..N"
    order = first_appearance(body, n_refs)
    orphans = sorted(set(range(1, n_refs + 1)) - set(order))
    if orphans:
        print(f"ERROR: never cited in main text: {orphans}")
        return 1
    mapping = {old: new for new, old in enumerate(order, start=1)}
    changed = {o: n for o, n in mapping.items() if o != n}
    print(f"{n_refs} references; {len(changed)} renumbered" + (":" if changed else "."))
    for o in sorted(changed):
        print(f"  {o:3d} -> {mapping[o]:3d}   {entries[o - 1][1][:70]}")

    outputs: dict[Path, str] = {}
    new_body, c, sus = rewrite(body, mapping, n_refs, "main")
    print(f"main.md: {c} citation groups" + (f"; left alone: {sus}" if sus else ""))
    by_new = sorted(entries, key=lambda e: mapping[e[0]])
    new_refs = "\n".join(f"{mapping[o]}. {t}" for o, t in by_new)
    outputs[MAIN] = f"{new_body}{REF_HEADER}\n\n{new_refs}\n"
    text, c, sus = rewrite(SI.read_text(), mapping, n_refs, SI.name)
    print(f"{SI.name}: {c} citation groups" + (f"; left alone: {sus}" if sus else ""))
    outputs[SI] = text
    # build.py is Python: only its FIGURES caption block may carry citations, so rewrite that slice
    # alone — tuples like (0, 1) elsewhere in the code would otherwise look like citations.
    btext = BUILD.read_text()
    head = re.search(r"^FIGURES\b[^\n]*\{\s*$", btext, re.M)
    if head is None:
        print(f"{BUILD.name}: no FIGURES block found; skipped")
    else:
        start = head.start()
        end = btext.index("\n}\n", start) + 3
        block, c, sus = rewrite(btext[start:end], mapping, n_refs, BUILD.name)
        print(f"{BUILD.name} captions: {c} citation groups" + (f"; left alone: {sus}" if sus else ""))
        outputs[BUILD] = btext[:start] + block + btext[end:]

    if apply:
        for path, text in outputs.items():
            path.write_text(text)
        print("applied.")
    else:
        print("dry run — pass --apply to write.")
    return 0


if __name__ == "__main__":
    sys.exit(main("--apply" in sys.argv))