#!/usr/bin/env python3
"""Linter: find numbered entities (equations, tables, figures) that are never
referenced in the text. Works on LaTeX sources AND compiled PDFs.

Guideline being enforced (ml-theses.org): every numbered equation, table, and
figure must be referenced in the text ("only number equations that are
referenced").

LaTeX mode (*.tex files or directories):
  1. UNREFERENCED    - \\label in a numbered environment that no \\ref, \\eqref,
                       \\autoref, \\cref/\\Cref, \\vref, \\pageref, or
                       \\hyperref points to.
  2. UNLABELED-EQ    - numbered math environment with no \\label at all.
  3. UNLABELED-FLOAT - captioned figure/table without a \\label.

PDF mode (*.pdf files) works on extracted text (pdftotext -layout, falling
back to PyMuPDF):
  Definitions detected:
    - caption lines starting with "Figure N:", "Table N:", "Algorithm N." etc.
    - equation numbers "(N)" right-aligned at the end of a line.
  Mentions detected:
    - "Figure N", "Fig. N", "Tables N and M", "Tables 11-17" (ranges/lists),
    - "Equation N", "Eq. (N)", and bare inline "(N)" (disable the bare form
      with --no-bare-eq-refs; bare "(N)" also matches enumerations, so keeping
      it ON avoids false alarms but may miss some unreferenced equations).
  An entity whose number never occurs outside its own caption/definition line
  is reported as UNREFERENCED. Lines with dot leaders (table of contents) are
  ignored.

PDF-mode limitations: heuristic text extraction; equation numbers only found
when right-aligned with >=2 spaces before "(N)"; mentions split across a page
break in mid-phrase may be missed (adjacent-line joins, including a hyphenated
word break such as "Ta-\nble 2.2", are handled).

Usage:
  python3 unreferenced_entity_linter.py thesis.pdf
  python3 unreferenced_entity_linter.py main.tex chapters/ [-a]
Exit status: 0 clean, 1 findings, 2 usage error.
"""

import argparse
import re
import subprocess
import sys
from pathlib import Path

# --------------------------------------------------------------------------
# LaTeX mode
# --------------------------------------------------------------------------

NUMBERED_MATH_ENVS = {
    "equation", "align", "gather", "multline", "eqnarray", "alignat", "flalign",
}
FLOAT_ENVS = {"figure", "table", "figure*", "table*", "algorithm", "listing"}

BEGIN_RE = re.compile(r"\\begin\{([A-Za-z*]+)\}")
END_RE_TMPL = r"\\end\{%s\}"
LABEL_RE = re.compile(r"\\label\{([^}]+)\}")
CAPTION_CMD_RE = re.compile(r"\\caption\b")
REF_RE = re.compile(r"\\(?:ref|eqref|autoref|cref|Cref|vref|pageref)\*?\{([^}]+)\}")
HYPERREF_RE = re.compile(r"\\hyperref\[([^\]]+)\]")
COMMENT_RE = re.compile(r"(?<!\\)%.*")


def strip_comments(line: str) -> str:
    return COMMENT_RE.sub("", line)


def classify(env, label):
    """Map a numbered entity to a kind (equation/figure/table/...).

    Prefers the enclosing environment name; when that is inconclusive,
    falls back to the ``prefix:`` part of the label (e.g. ``fig:flow``).
    """
    base = env.rstrip("*") if env else ""
    if base in NUMBERED_MATH_ENVS:
        return "equation"
    if base in {"figure", "table", "algorithm", "listing"}:
        return base
    prefix = label.split(":", 1)[0].lower() if ":" in label else ""
    return {
        "eq": "equation", "fig": "figure", "tab": "table", "table": "table",
        "alg": "algorithm", "sec": "section", "ch": "chapter", "chap": "chapter",
        "app": "appendix", "lst": "listing", "thm": "theorem", "lem": "lemma",
    }.get(prefix, "other")


def scan_tex_file(path: Path):
    """Scan one .tex file for numbered entities and references.

    Returns ``(entities, refs, unlabeled_math, unlabeled_floats)``:
    ``entities`` are ``\\label`` sites (with kind/env/location), ``refs`` is
    the set of all referenced label names, and the two ``unlabeled_*`` lists
    hold numbered math / captioned floats that carry no ``\\label``. A stack
    tracks nested numbered environments so a label/caption is attributed to
    the innermost one.
    """
    lines = [strip_comments(l) for l in
             path.read_text(encoding="utf-8", errors="replace").splitlines()]
    entities, refs = [], set()
    unlabeled_math, unlabeled_floats = [], []

    for ln in lines:
        for m in REF_RE.finditer(ln):
            refs.update(x.strip() for x in m.group(1).split(","))
        for m in HYPERREF_RE.finditer(ln):
            refs.add(m.group(1).strip())

    stack = []  # [env, start_line, has_label, has_caption, starred]
    for lineno, ln in enumerate(lines, start=1):
        pos = 0
        while True:
            b = BEGIN_RE.search(ln, pos)
            e = None
            for entry in stack:
                m = re.search(END_RE_TMPL % re.escape(entry[0]), ln[pos:])
                if m and (e is None or m.start() < e[0].start()):
                    e = (m, entry)
            if b and (e is None or b.start() - pos < e[0].start()):
                env = b.group(1)
                if env.rstrip("*") in NUMBERED_MATH_ENVS or env in FLOAT_ENVS:
                    stack.append([env, lineno, False, False, env.endswith("*")])
                pos = b.end()
                continue
            if e:
                m, entry = e
                env, start, has_label, has_caption, starred = entry
                if env.rstrip("*") in NUMBERED_MATH_ENVS and not starred and not has_label:
                    unlabeled_math.append({"env": env, "file": path, "line": start})
                if env in FLOAT_ENVS and has_caption and not has_label:
                    unlabeled_floats.append({"env": env, "file": path, "line": start})
                stack.remove(entry)
                pos += m.end()
                continue
            break

        for m in LABEL_RE.finditer(ln):
            label = m.group(1).strip()
            env = stack[-1][0] if stack else None
            if stack:
                stack[-1][2] = True
            entities.append({"label": label, "kind": classify(env, label),
                             "env": env, "file": path, "line": lineno})
        if stack and CAPTION_CMD_RE.search(ln):
            stack[-1][3] = True

    return entities, refs, unlabeled_math, unlabeled_floats


def run_tex_mode(files, all_labels):
    """Run the LaTeX-mode check across all .tex files and print findings.

    Aggregates every file's entities/refs, then reports each labeled entity
    whose label is never referenced (restricted to float/equation kinds
    unless ``all_labels``), plus the UNLABELED-EQ / UNLABELED-FLOAT cases.
    Returns the total number of findings.
    """
    all_entities, all_refs = [], set()
    all_umath, all_ufloat = [], []
    for f in files:
        ents, refs, umath, ufloat = scan_tex_file(f)
        all_entities.extend(ents)
        all_refs |= refs
        all_umath.extend(umath)
        all_ufloat.extend(ufloat)

    checked = {"equation", "figure", "table", "algorithm", "listing"}
    findings = 0
    for ent in sorted((e for e in all_entities
                       if e["label"] not in all_refs
                       and (all_labels or e["kind"] in checked)),
                      key=lambda x: (str(x["file"]), x["line"])):
        print(f"{ent['file']}:{ent['line']}: UNREFERENCED {ent['kind']} "
              f"'\\label{{{ent['label']}}}' is never referenced in the text")
        findings += 1
    for it in all_umath:
        print(f"{it['file']}:{it['line']}: UNLABELED-EQ numbered '{it['env']}' "
              f"environment has no \\label -- its number can never be "
              f"referenced (star it, or label and reference it)")
        findings += 1
    for it in all_ufloat:
        print(f"{it['file']}:{it['line']}: UNLABELED-FLOAT captioned "
              f"'{it['env']}' has no \\label -- numbered but unreferenceable")
        findings += 1

    n = sum(1 for e in all_entities if all_labels or e["kind"] in checked)
    print(f"\n[tex] checked {len(files)} file(s), {n} labeled entit(ies), "
          f"{len(all_refs)} referenced label(s): {findings} finding(s)")
    return findings


# --------------------------------------------------------------------------
# PDF mode
# --------------------------------------------------------------------------

FLOAT_WORDS = {
    "figure": "figure", "figures": "figure", "fig.": "figure", "figs.": "figure",
    "table": "table", "tables": "table",
    "algorithm": "algorithm", "algorithms": "algorithm",
    "listing": "listing", "listings": "listing",
}
NUM = r"[A-Z]?\d+(?:\.\d+)?"
# Separator after the number must be a colon, or a period FOLLOWED by
# whitespace/end -- never the period inside the number itself. Plain "[:.]"
# let the regex backtrack "2.4" to "2" and match the internal dot, recording a
# phantom "Figure 2" definition from a body line like "Figure 2.4 summarises
# these layers:" and then reporting that non-existent float as unreferenced.
CAPTION_DEF_RE = re.compile(
    rf"^\s*(Figure|Fig\.|Table|Algorithm|Listing)\s+({NUM})\s*(?::|\.(?=\s|$))")
EQ_DEF_RE = re.compile(rf"(?:\s{{2,}}|^\s*)\(({NUM})\)\s*$")
FLOAT_MENTION_RE = re.compile(
    rf"\b(Figures?|Figs?\.|Tables?|Algorithms?|Listings?)\s+"
    rf"({NUM}(?:\s*(?:,|and|&|/|to|--?|–|—)\s*{NUM})*)")
EQ_WORD_MENTION_RE = re.compile(
    rf"\b(?:Equations?|Eqs?\.)\s*\(?({NUM}(?:\)?\s*(?:,|and|&|--?|–)\s*\(?{NUM})*)\)?")
BARE_EQ_RE = re.compile(r"\((\d{1,3})\)")
DOT_LEADER_RE = re.compile(r"(?:\.\s){4,}|\.{4,}")
NUM_TOKEN_RE = re.compile(NUM)


def extract_pdf_pages(path: Path):
    """Return the PDF's text as a list of per-page strings.

    Uses ``pdftotext -layout`` (poppler) when available, otherwise falls
    back to PyMuPDF; exits with an error if neither can be used.
    """
    try:
        out = subprocess.run(["pdftotext", "-layout", str(path), "-"],
                             capture_output=True, check=True)
        return out.stdout.decode("utf-8", errors="replace").split("\f")
    except (FileNotFoundError, subprocess.CalledProcessError):
        pass
    try:
        import fitz  # PyMuPDF
    except ImportError:
        sys.exit("error: need either pdftotext (poppler) or PyMuPDF to read PDFs")
    with fitz.open(path) as doc:
        return [page.get_text() for page in doc]


def expand_number_list(list_str):
    """'11-17' -> 11..17; '5 and 6' -> 5,6; '13/16/19' -> 13,16,19."""
    nums = NUM_TOKEN_RE.findall(list_str)
    out = list(nums)
    for m in re.finditer(rf"({NUM})\s*(?:--?|–|—|to)\s*({NUM})", list_str):
        a, b = m.group(1), m.group(2)
        if a.isdigit() and b.isdigit() and int(a) < int(b) <= int(a) + 50:
            out.extend(str(i) for i in range(int(a) + 1, int(b)))
    return out


def sort_key(entity):
    """Order a ``(kind, number)`` entity numerically, not lexically, so that
    e.g. "Figure 2" sorts before "Figure 10" and "A.1" groups by prefix."""
    kind, num = entity
    m = re.match(r"([A-Z]*)(\d+)(?:\.(\d+))?", num)
    return (kind, m.group(1), int(m.group(2)), int(m.group(3) or 0))


def scan_pdf(path: Path, bare_eq_refs=True):
    """Run the PDF-mode check on one file and print findings.

    Collects entity *definitions* (caption lines, right-aligned equation
    numbers) and *mentions* (float/equation references, optionally bare
    ``(N)``) from the extracted text, then reports each defined entity whose
    number is never mentioned outside its own definition. Returns the number
    of findings. See the module docstring for the detection heuristics.
    """
    pages = extract_pdf_pages(path)
    defs = {}      # (kind, number) -> first page of definition
    mentions = set()  # (kind, number)

    for pageno, page in enumerate(pages, start=1):
        lines = page.splitlines()
        for i, line in enumerate(lines):
            if DOT_LEADER_RE.search(line):
                continue  # table-of-contents / list-of-figures entry

            cap = CAPTION_DEF_RE.match(line)
            if cap:
                kind = FLOAT_WORDS[cap.group(1).lower()]
                defs.setdefault((kind, cap.group(2)), pageno)

            eq = EQ_DEF_RE.search(line)
            if eq:
                defs.setdefault(("equation", eq.group(1)), pageno)

            for m in FLOAT_MENTION_RE.finditer(line):
                # skip the caption's own "Figure N:" occurrence
                if cap and m.start(1) == cap.start(1):
                    continue
                kind = FLOAT_WORDS[m.group(1).lower()]
                for n in expand_number_list(m.group(2)):
                    mentions.add((kind, n))

            for m in EQ_WORD_MENTION_RE.finditer(line):
                for n in expand_number_list(m.group(1)):
                    mentions.add(("equation", n))

            if bare_eq_refs:
                bares = list(BARE_EQ_RE.finditer(line))
                if eq and bares and bares[-1].start() >= eq.start():
                    bares = bares[:-1]  # last one is the definition itself
                for m in bares:
                    mentions.add(("equation", m.group(1)))

            # A mention split across a line break, including a HYPHENATED
            # word break: "... Table\n2.2 ..." or "... Ta-\nble 2.2 ...".
            # Heal the seam (join the hyphenated word, else glue the last word
            # to the next line) and re-scan it, so the reference is still seen.
            # Additive: this can only add mentions, never suppress a finding.
            if i + 1 < len(lines):
                cur = line.rstrip()
                nxt = lines[i + 1].lstrip()
                seam = (cur[:-1].rsplit(" ", 1)[-1] + nxt) if cur.endswith("-") \
                    else (cur.rsplit(" ", 1)[-1] + " " + nxt)
                for m in FLOAT_MENTION_RE.finditer(seam):
                    kind = FLOAT_WORDS[m.group(1).lower()]
                    for n in expand_number_list(m.group(2)):
                        mentions.add((kind, n))
                for m in EQ_WORD_MENTION_RE.finditer(seam):
                    for n in expand_number_list(m.group(1)):
                        mentions.add(("equation", n))

    findings = 0
    for (kind, num), page in sorted(defs.items(), key=lambda kv: sort_key(kv[0])):
        if (kind, num) not in mentions:
            label = f"({num})" if kind == "equation" else f"{kind.capitalize()} {num}"
            print(f"{path}: UNREFERENCED {kind} {label} "
                  f"(defined on PDF page {page}) is never mentioned in the text")
            findings += 1

    n_eq = sum(1 for k, _ in defs if k == "equation")
    print(f"\n[pdf] {path.name}: {len(defs)} numbered entit(ies) found "
          f"({n_eq} equations, {len(defs) - n_eq} floats), "
          f"{findings} finding(s)"
          + ("" if bare_eq_refs else "  [bare '(N)' not counted as references]"))
    return findings


# --------------------------------------------------------------------------

def main(argv=None):
    ap = argparse.ArgumentParser(
        description="Find numbered equations/tables/figures never referenced "
                    "in the text (LaTeX sources or PDF).")
    ap.add_argument("paths", nargs="+", help=".tex files, directories, or .pdf files")
    ap.add_argument("-a", "--all-labels", action="store_true",
                    help="[tex] also report unreferenced non-float labels")
    ap.add_argument("--no-bare-eq-refs", action="store_true",
                    help="[pdf] do not count bare inline '(N)' as an equation "
                         "reference (stricter; may flag enumerations)")
    args = ap.parse_args(argv)

    tex_files, pdf_files = [], []
    for p in map(Path, args.paths):
        if p.is_dir():
            tex_files.extend(sorted(p.rglob("*.tex")))
        elif p.is_file() and p.suffix.lower() == ".pdf":
            pdf_files.append(p)
        elif p.is_file():
            tex_files.append(p)
        else:
            sys.exit(f"error: no such file or directory: {p}")
    if not tex_files and not pdf_files:
        sys.exit("error: no .tex or .pdf files found")

    print("Checks that every numbered equation, table, and figure is "
          "referred to at least once in the text (and, in LaTeX mode, that "
          "each has a \\label). Each line below is one unreferenced entity.\n")

    findings = 0
    if tex_files:
        findings += run_tex_mode(tex_files, args.all_labels)
    for pdf in pdf_files:
        findings += scan_pdf(pdf, bare_eq_refs=not args.no_bare_eq_refs)
    return 1 if findings else 0


if __name__ == "__main__":
    sys.exit(main())