arborist/tests/test_concepts_extract.py

"""Tests for ``arborist.concepts.extract`` — concept-relation
extractors that derive synonyms / IDF / FTS5-titles from corpus
state. Module had zero direct tests despite being the write-side
of every concept-relations row the existing test_concepts.py
exercises on the read-side.

Coverage:
  - _title_tokens (pure private helper) — stopword strip,
    underscore→space, length≥4 filter, lowercase, hyphen-handling
  - EXTRACTORS registry contract — keys + signature
  - backfill_documents_fts on a synthetic mini-shard
  - link_reciprocity_synonym idempotency on a 2-doc reciprocal pair

The extractors operate on real shard schema; we build a minimal
in-memory schema that matches store.py's contract just enough to
exercise the code paths.
"""

from __future__ import annotations

import sqlite3

import pytest

from arborist.concepts.extract import (
    EXTRACTORS,
    _title_tokens,
    backfill_documents_fts,
    backfill_token_idf,
    link_reciprocity_synonym,
)
from arborist.store import SCHEMA_SQL


# --- _title_tokens (pure) ------------------------------------------


def test_title_tokens_empty_string():
    assert _title_tokens("") == set()


def test_title_tokens_strips_stopwords():
    """Documented stopwords (the/a/an/of/and/...) drop out."""
    out = _title_tokens("the history of computing")
    assert "the" not in out  # stopword
    assert "history" in out
    assert "computing" in out
    # "of" stopword + length<4 anyway
    assert "of" not in out


def test_title_tokens_underscore_to_space():
    """Wikipedia title style 'Foo_Bar' → split on underscore."""
    out = _title_tokens("New_York_City")
    assert "york" in out
    assert "city" in out


def test_title_tokens_length_four_floor():
    """Tokens length<4 dropped (3-letter words = noisy anchors)."""
    out = _title_tokens("the cat ran")
    assert "cat" not in out  # 3 chars
    assert "ran" not in out  # 3 chars


def test_title_tokens_lowercase():
    """Output always lowercase regardless of input case."""
    out = _title_tokens("Quantum CHROMODYNAMICS")
    assert "quantum" in out
    assert "chromodynamics" in out


def test_title_tokens_handles_hyphenated_token():
    """Title-token regex `[a-z][a-z0-9'\\-]+` keeps hyphens inside
    a token. 'state-of-the-art' → 'state-of-the-art' as one token
    (length>=4)."""
    out = _title_tokens("state-of-the-art technology")
    # The regex starts on a letter then captures alphanum/apos/hyphen.
    # 'state-of-the-art' becomes one token.
    assert any("-" in t for t in out)
    assert "technology" in out


def test_title_tokens_dedupe():
    """Output is a set — duplicates collapse."""
    out = _title_tokens("Quantum quantum mechanics")
    quantum_count = sum(1 for t in out if t == "quantum")
    assert quantum_count == 1


def test_title_tokens_skips_pure_digits():
    """Token must START with [a-z] per the regex; 1984 → no match."""
    out = _title_tokens("1984 novel")
    assert "1984" not in out
    assert "novel" in out


# --- EXTRACTORS registry --------------------------------------------


def test_extractors_registry_has_documented_keys():
    """Three extractors documented in the module-level docstring."""
    assert "link_reciprocity" in EXTRACTORS
    assert "token_idf" in EXTRACTORS
    assert "documents_fts" in EXTRACTORS


def test_extractors_registry_callables():
    """Each registry value is callable."""
    for name, fn in EXTRACTORS.items():
        assert callable(fn), f"EXTRACTORS[{name!r}] is not callable"


def test_extractors_registry_link_reciprocity_routes_to_function():
    """The dispatch table maps to the actual implementations."""
    assert EXTRACTORS["link_reciprocity"] is link_reciprocity_synonym
    assert EXTRACTORS["token_idf"] is backfill_token_idf
    assert EXTRACTORS["documents_fts"] is backfill_documents_fts


# --- backfill_documents_fts ----------------------------------------


@pytest.fixture
def empty_shard(tmp_path):
    """Fresh SQLite with arborist's full schema applied."""
    db_path = tmp_path / "shard.db"
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_SQL)
    yield conn
    conn.close()


def _insert_doc(conn, root: str, uri: str, title: str):
    conn.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, "
        " compression_depth, title, chunking_version, "
        " canonicalization_version, schema_version, ingest_ts) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (root, uri, "html", "surface", 0, title,
         "tok-512-v1", "norm-v1", "v9.8.0", 0),
    )


def test_backfill_documents_fts_indexes_titles(empty_shard):
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Apple Inc")
    _insert_doc(empty_shard, "b" * 64, "uri-b", "Banana fruit")
    _insert_doc(empty_shard, "c" * 64, "uri-c", "Cherry tree")
    empty_shard.commit()

    result = backfill_documents_fts(empty_shard)
    assert result["rows_indexed"] == 3
    assert result["elapsed_ms"] >= 0


def test_backfill_documents_fts_idempotent(empty_shard):
    """Re-running on the same shard yields identical row count
    (DELETE + INSERT inside)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "First")
    _insert_doc(empty_shard, "b" * 64, "uri-b", "Second")
    empty_shard.commit()

    r1 = backfill_documents_fts(empty_shard)
    r2 = backfill_documents_fts(empty_shard)
    assert r1["rows_indexed"] == r2["rows_indexed"] == 2


def test_backfill_documents_fts_skips_null_titles(empty_shard):
    """Documents with NULL title are not indexed (no FTS5 row)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Has Title")
    # Insert a row with no title.
    empty_shard.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, "
        " compression_depth, title, chunking_version, "
        " canonicalization_version, schema_version, ingest_ts) "
        "VALUES (?, ?, ?, ?, ?, NULL, ?, ?, ?, ?)",
        ("b" * 64, "uri-b", "html", "surface", 0,
         "tok-512-v1", "norm-v1", "v9.8.0", 0),
    )
    empty_shard.commit()
    result = backfill_documents_fts(empty_shard)
    assert result["rows_indexed"] == 1  # only the titled doc


def test_backfill_documents_fts_default_derived_from_doesnt_break(empty_shard):
    """Default derived_from kwarg path works (None → derived from
    the function name internally)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Test")
    empty_shard.commit()
    # Default kwarg path
    r1 = backfill_documents_fts(empty_shard)
    # Explicit kwarg path
    r2 = backfill_documents_fts(empty_shard, derived_from="custom-extractor")
    assert r1["rows_indexed"] == r2["rows_indexed"]


# --- link_reciprocity_synonym --------------------------------------


def _insert_edge(conn, src: str, dst: str, dst_root: str):
    conn.execute(
        "INSERT INTO edges "
        "(src_root, edge_type, dst_root, dst_uri, anchor) "
        "VALUES (?, 'related_doc', ?, ?, '')",
        (src, dst_root, dst),
    )


def test_link_reciprocity_no_edges_returns_zero(empty_shard):
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 0
    assert result["synonyms_inserted"] == 0


def test_link_reciprocity_one_way_edge_no_pair(empty_shard):
    """Edge A→B without reverse B→A → no reciprocal pair → no
    synonym."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple Inc")
    _insert_doc(empty_shard, b, "uri-b", "Banana fruit")
    _insert_edge(empty_shard, a, "uri-b", b)  # only A→B
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 0


def test_link_reciprocity_pair_creates_synonyms(empty_shard):
    """A↔B with title tokens → synonym edges."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana technology")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 1
    # Cross-product of 2 tokens × 2 tokens, minus self-overlap.
    # apple/computing × banana/technology — no overlap → 4 pairs.
    assert result["synonyms_inserted"] == 4


def test_link_reciprocity_idempotent(empty_shard):
    """Re-running adds no new rows (concept_relations PK enforces
    uniqueness)."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana technology")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    r1 = link_reciprocity_synonym(empty_shard)
    r2 = link_reciprocity_synonym(empty_shard)
    assert r1["synonyms_inserted"] >= 1
    assert r2["synonyms_skipped"] == r1["synonyms_inserted"]
    assert r2["synonyms_inserted"] == 0


def test_link_reciprocity_skips_self_overlap_token(empty_shard):
    """If both titles share a token (e.g. 'computing' in both),
    that's not a synonym (same word) — skipped from the cross
    product."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana computing")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    # 4 cross pairs minus the (computing, computing) self-overlap = 3
    # apple→banana, apple→computing, banana→computing, computing→apple,
    # computing→banana. Self-pair (computing, computing) excluded.
    # Implementation: cross = {(a,b) for a in A for b in B if a != b}
    # A = {apple, computing}, B = {banana, computing}
    # → (apple, banana), (apple, computing), (computing, banana) = 3
    assert result["synonyms_inserted"] == 3


# --- acronym_parens_synonym (#000054) -------------------------------


def _insert_lead_chunk(conn, doc_root: str, text: str):
    """Insert a single idx=0 chunk for ``doc_root`` carrying ``text``
    in the schema-expected packed form."""
    from arborist.compress import pack_chunk
    content = pack_chunk(text)
    conn.execute(
        "INSERT INTO chunks "
        "(document_root, idx, leaf_hash, content, tier) "
        "VALUES (?, ?, ?, ?, ?)",
        (doc_root, 0, "0" * 64, content, "hot"),
    )


def _doc_with_lead(conn, root: str, uri: str, title: str, body: str):
    _insert_doc(conn, root, uri, title)
    _insert_lead_chunk(conn, root, body)


def test_acronym_parens_emits_bidirectional_synonym_on_cpu_case(empty_shard):
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard,
        "a" * 64, "uri-a", "Central processing unit",
        "A central processing unit (CPU) is the most important component "
        "of a computer.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 1
    # 3 content words (central, processing, unit) × 2 directions = 6 edges
    assert r["synonyms_inserted"] == 6
    rows = empty_shard.execute(
        "SELECT token, target FROM concept_relations "
        "WHERE evidence_kind = 'acronym_parens' ORDER BY token, target"
    ).fetchall()
    pairs = {(r["token"], r["target"]) for r in rows}
    for word in ("central", "processing", "unit"):
        assert ("cpu", word) in pairs
        assert (word, "cpu") in pairs


def test_acronym_parens_idempotent(empty_shard):
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "RAM",
        "Random Access Memory (RAM) is volatile.",
    )
    empty_shard.commit()
    r1 = acronym_parens_synonym(empty_shard)
    r2 = acronym_parens_synonym(empty_shard)
    assert r1["synonyms_inserted"] >= 1
    assert r2["synonyms_inserted"] == 0
    assert r2["synonyms_skipped"] == r1["synonyms_inserted"]


def test_acronym_parens_filters_function_words_in_phrase(empty_shard):
    """'Federal Bureau of Investigation (FBI)' — 'of' is dropped before
    initial matching, leaving F/B/I to match Federal/Bureau/Investigation."""
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "FBI",
        "The Federal Bureau of Investigation (FBI) is a US agency.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 1
    rows = empty_shard.execute(
        "SELECT token, target FROM concept_relations "
        "WHERE token = 'fbi'"
    ).fetchall()
    targets = {row["target"] for row in rows}
    assert {"federal", "bureau", "investigation"} <= targets
    assert "of" not in targets  # function word — never an edge


def test_acronym_parens_rejects_mismatched_initials(empty_shard):
    """'Hypertext Transfer Protocol (HTTP)' — 4-letter acronym but only
    3 content words → strict 1:1 fails → no edge emitted."""
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "HTTP",
        "Hypertext Transfer Protocol (HTTP) is a network protocol.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 0
    assert r["synonyms_inserted"] == 0


def test_acronym_parens_rejects_initials_mismatch_at_position(empty_shard):
    """'Apple Banana Carrot (XYZ)' — letters don't match in order → no edge."""
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "Misleading",
        "Apple Banana Carrot (XYZ) is a fake expansion.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 0


def test_acronym_parens_handles_hyphenated_words(empty_shard):
    """'Read-only memory (ROM)' — 'Read-only' splits into Read+only;
    'only' is a function-style word but NOT in skipwords, so atoms =
    [Read, only, memory] → ROM (R/O/M) → match."""
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "ROM",
        "Read-only memory (ROM) cannot be written.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 1
    rows = empty_shard.execute(
        "SELECT token, target FROM concept_relations WHERE token = 'rom'"
    ).fetchall()
    targets = {row["target"] for row in rows}
    assert {"read", "memory"} <= targets  # "only" is 4 chars, kept; "read" 4 chars
    # 'only' is included as a target since len >= 3
    assert "only" in targets


def test_acronym_parens_dedupes_repeated_definition_in_same_doc(empty_shard):
    """Same parenthetical pattern repeated → one edge set per doc."""
    from arborist.concepts.extract import acronym_parens_synonym
    _doc_with_lead(
        empty_shard, "a" * 64, "uri-a", "CPU",
        "A central processing unit (CPU) is the brain. "
        "The central processing unit (CPU) executes instructions.",
    )
    empty_shard.commit()
    r = acronym_parens_synonym(empty_shard)
    assert r["pairs_found"] == 1
    assert r["synonyms_inserted"] == 6  # 3 words × 2 directions


def test_acronym_parens_in_extractors_registry():
    from arborist.concepts.extract import EXTRACTORS
    assert "acronym_parens" in EXTRACTORS
    assert callable(EXTRACTORS["acronym_parens"])