tests/test_lexical_refinement.py

"""Conservative lexical refinement and rollback tests."""

from __future__ import annotations

from dataclasses import dataclass

import app.engine.lexical as lexical
import app.engine.orchestrator as orchestrator
from app.engine.lexical import LexicalResult, refine_sentence
from app.engine.models import LexicalChange


@dataclass
class FakeWord:
    value: str

    def lemma(self) -> str:
        return self.value


class FakeSynset:
    def __init__(
        self,
        synset_id: str,
        definition: str,
        words: list[str],
        examples: list[str] | None = None,
    ):
        self.id = synset_id
        self._definition = definition
        self._words = [FakeWord(word) for word in words]
        self._examples = examples or []

    def definition(self) -> str:
        return self._definition

    def examples(self) -> list[str]:
        return self._examples

    def words(self) -> list[FakeWord]:
        return self._words


class FakeWordnet:
    def __init__(self, mapping: dict[tuple[str, str], list[FakeSynset]]):
        self.mapping = mapping

    def synsets(self, lemma: str, *, pos: str) -> list[FakeSynset]:
        return self.mapping.get((lemma, pos), [])


def provide_wordnet() -> FakeWordnet:
    """Fake lexicon where candidates are everyday / more common words."""
    return FakeWordnet(
        {
            ("assist", "v"): [
                FakeSynset(
                    "fake-assist-v",
                    "help students with useful information",
                    ["assist", "help"],
                )
            ],
            ("vital", "a"): [
                FakeSynset(
                    "fake-vital-a",
                    "needed detailed information that helps readers",
                    ["vital", "needed"],
                )
            ],
        }
    )


def test_context_supported_candidate_prefers_simpler_word():
    result = refine_sentence(
        "The report assists students with useful information.",
        min_wsd=0.18,
        wordnet=provide_wordnet(),
    )
    assert result.text == "The report helps students with useful information."
    assert len(result.changes) == 1
    assert result.changes[0].original == "assists"
    assert result.changes[0].replacement == "helps"
    assert result.changes[0].synset_id == "fake-assist-v"


def test_no_gloss_overlap_means_no_change():
    resource = FakeWordnet(
        {
            ("assist", "v"): [
                FakeSynset(
                    "fake-unrelated-v1",
                    "equip a room with electrical machinery",
                    ["assist", "abet"],
                ),
                FakeSynset(
                    "fake-unrelated-v2",
                    "operate a machine in a factory hall",
                    ["assist", "abet"],
                ),
            ]
        }
    )
    source = "The report assists students with useful information."
    result = refine_sentence(source, min_wsd=0.18, wordnet=resource)
    assert result.text == source
    assert not result.changes


def test_harder_advanced_synonym_is_rejected():
    resource = FakeWordnet(
        {
            ("achieve", "v"): [
                FakeSynset(
                    "fake-achieve-v",
                    "to gain with effort",
                    ["achieve", "accomplish"],
                )
            ]
        }
    )
    source = "People achieve goals every year."
    result = refine_sentence(source, min_wsd=0.18, wordnet=resource)
    # "accomplish" is less common than "achieve" — keep the simpler source.
    assert result.text == source
    assert not result.changes


def test_simpler_everyday_synonym_is_preferred():
    resource = FakeWordnet(
        {
            ("assist", "v"): [
                FakeSynset(
                    "fake-assist-v",
                    "help people finish work",
                    ["assist", "help"],
                )
            ]
        }
    )
    result = refine_sentence(
        "Teachers assist people every day.",
        min_wsd=0.18,
        wordnet=resource,
    )
    assert result.text == "Teachers help people every day."
    assert result.changes[0].replacement == "help"


def test_article_agrees_after_simpler_adjective():
    resource = FakeWordnet(
        {
            ("significant", "a"): [
                FakeSynset(
                    "fake-significant-a",
                    "important problem for people",
                    ["significant", "important"],
                )
            ]
        }
    )
    result = refine_sentence(
        "This is a significant problem for people.",
        min_wsd=0.18,
        wordnet=resource,
    )
    assert "an important problem" in result.text
    assert "a important" not in result.text


def test_skill_is_not_swapped_to_science():
    source = (
        "Developing teamwork skills not only improves project outcomes "
        "but also prepares individuals for future career opportunities."
    )
    result = refine_sentence(source, min_wsd=0.14, max_changes=3)
    assert "sciences" not in result.text.lower()
    assert "skills" in result.text.lower()
    assert "persons" not in result.text.lower()


def test_bad_teamwork_polish_patterns_are_rejected(monkeypatch):
    source = (
        "They can accomplish projects more efficiently and produce better "
        "results when people collaborate effectively."
    )
    monkeypatch.setattr(
        orchestrator,
        "paraphrase_sentence",
        lambda text, **_kwargs: type(
            "R",
            (),
            {
                "text": (
                    "They can action projects more efficiently and produce "
                    "better results when people effectively collaborate."
                ),
                "confidence": 0.9,
            },
        )(),
    )
    result = orchestrator.rewrite_document(
        source,
        force_rewrite=False,
        use_paraphrase=True,
        use_lexical_refinement=True,
        use_minilm_safety=False,
    )
    assert "action projects" not in result.text.lower()
    assert "accomplish" in result.text.lower() or "when people" in result.text.lower()


def test_teamwork_sample_reorders_without_bad_synonyms():
    source = (
        "Teamwork is an essential skill in both academic and professional "
        "environments. They can accomplish projects more efficiently and "
        "produce better results when people collaborate effectively. Unique "
        "skills, experiences, and perspectives that bring to working problems "
        "creatively are brought by each team member."
    )
    result = orchestrator.rewrite_document(
        source,
        force_rewrite=False,
        use_paraphrase=False,
        use_lexical_refinement=True,
        use_minilm_safety=False,
    )
    text = result.text.lower()
    assert "action projects" not in text
    assert "positions that" not in text
    assert "get to working" not in text
    assert "team work" not in text
    # At least the first two clear sentences should change via structure.
    changed = [
        s
        for s in result.sentences
        if s.original.strip()
        and s.rewritten.strip().lower().rstrip(".!?")
        != s.original.strip().lower().rstrip(".!?")
    ]
    assert len(changed) >= 2


def test_equipped_to_and_teamwork_stay_natural():
    source = (
        "Good communication is the foundation of successful teamwork. "
        "Listening to others, sharing ideas respectfully, and supporting "
        "teammates help build trust and improve collaboration. "
        "Teams that work well together are often more productive and "
        "better equipped to overcome challenges."
    )
    result = orchestrator.rewrite_document(
        source,
        force_rewrite=False,
        use_paraphrase=False,
        use_lexical_refinement=True,
        use_minilm_safety=False,
        require_wording_change=False,
    )
    text = result.text.lower()
    assert "team work" not in text
    assert "teamwork" in text
    assert "fitted" not in text
    assert "equipped" in text
    assert "master challenges" not in text
    assert "supporting mates" not in text
    assert "teammates" in text


def test_rare_synset_alternative_is_rejected():
    resource = FakeWordnet(
        {
            ("procrastination", "n"): [
                FakeSynset(
                    "fake-procrastination-n",
                    "delaying an important action until later",
                    ["procrastination", "cunctation"],
                )
            ]
        }
    )
    source = "Procrastination means delaying an important action until later."
    result = refine_sentence(source, min_wsd=0.18, wordnet=resource)
    assert result.text == source
    assert not result.changes


def test_protected_and_quoted_text_is_excluded():
    resource = provide_wordnet()
    protected = "The report assists ZZPROTECTEDEMAIL0ZZ with useful information."
    assert not refine_sentence(
        protected,
        min_wsd=0.18,
        wordnet=resource,
    ).changes
    quoted = '"The report assists students with useful information."'
    assert not refine_sentence(
        quoted,
        min_wsd=0.18,
        wordnet=resource,
    ).changes
    citation = "The report assists students with useful information (Smith, 2024)."
    assert not refine_sentence(
        citation,
        min_wsd=0.18,
        wordnet=resource,
    ).changes


def test_max_changes_is_capped():
    result = refine_sentence(
        "The vital report assists students with detailed information for readers.",
        min_wsd=0.18,
        max_changes=99,
        wordnet=provide_wordnet(),
    )
    assert 1 <= len(result.changes) <= 15


def test_dynamic_budget_scales_with_sentence_length():
    from app.engine.lexical import dynamic_lexical_budget

    short = "Students need help."
    long = (
        "The research team carefully examined several important documents "
        "before presenting their detailed findings to the committee members "
        "during the lengthy afternoon session yesterday."
    )
    assert dynamic_lexical_budget(short) >= 1
    assert dynamic_lexical_budget(long) > dynamic_lexical_budget(short)
    assert dynamic_lexical_budget(long, polish=True) >= dynamic_lexical_budget(long)
    assert dynamic_lexical_budget(long) <= 15


def test_polish_can_add_extra_changes_even_on_short_sentence():
    result_plain = refine_sentence(
        "The manager subsequently assisted several diligent students during the unusually difficult afternoon workshop.",
        min_wsd=0.14,
        polish=False,
    )
    result_polish = refine_sentence(
        "The manager subsequently assisted several diligent students during the unusually difficult afternoon workshop.",
        min_wsd=0.14,
        polish=True,
    )
    assert len(result_plain.changes) >= 2
    assert len(result_polish.changes) > len(result_plain.changes)
    assert any(change.replacement == "later" for change in result_polish.changes)
    assert any(change.replacement == "hard" for change in result_polish.changes)


def test_polish_uses_related_adjective_senses():
    source = (
        "Effective communication is an essential skill that creates "
        "significant results in numerous workplaces."
    )
    plain = refine_sentence(source, min_wsd=0.14, polish=False)
    polished = refine_sentence(source, min_wsd=0.14, polish=True)
    assert len(polished.changes) > len(plain.changes)
    polished_text = polished.text.lower()
    assert "necessary" in polished_text or "useful" in polished_text or "efficient" in polished_text
    assert "hard communication" not in polished_text
    assert "big skill" not in polished_text
    assert "big results" not in polished_text


def test_identify_areas_is_not_named_by_polish():
    source = (
        "Businesses that actively listen to customer feedback can identify "
        "areas for improvement and strengthen customer loyalty."
    )
    polished = refine_sentence(source, min_wsd=0.14, polish=True)
    assert "name areas" not in polished.text.lower()
    assert "place areas" not in polished.text.lower()
    assert "identify" in polished.text.lower()


def test_satisfying_customer_experience_is_not_filling_or_taking():
    source = (
        "Employees who communicate clearly and maintain a positive attitude "
        "help create a satisfying customer experience."
    )
    plain = refine_sentence(source, min_wsd=0.14, polish=False)
    polished = refine_sentence(source, min_wsd=0.14, polish=True)
    for result in (plain, polished):
        low = result.text.lower()
        assert "filling customer" not in low
        assert "taking customer" not in low
        assert "satisfying" in low or "keep" in low
    assert len(polished.changes) > len(plain.changes)
    assert any(
        change.original.lower() == "maintain" and change.replacement.lower() == "keep"
        for change in polished.changes
    )


def test_customer_experience_rewrite_diverges_with_polish():
    source = (
        "Employees who communicate clearly and maintain a positive attitude "
        "help create a satisfying customer experience. Businesses that actively "
        "listen to customer feedback can identify areas for improvement and "
        "strengthen customer loyalty."
    )
    plain = orchestrator.rewrite_document(
        source,
        lexical_polish=False,
        use_lexical_refinement=True,
        use_paraphrase=True,
        use_minilm_safety=False,
    )
    polished = orchestrator.rewrite_document(
        source,
        lexical_polish=True,
        use_lexical_refinement=True,
        use_paraphrase=True,
        use_minilm_safety=False,
    )
    assert "filling customer" not in plain.text.lower()
    assert "taking customer" not in polished.text.lower()
    assert "satisfying" in plain.text.lower()
    assert plain.text != polished.text
    assert "keep" in polished.text.lower()


def test_reputation_retention_sentence_gets_polish_synonym():
    source = (
        "By consistently delivering high-quality service, organizations can "
        "build a strong reputation, increase customer retention, and "
        "encourage positive word-of-mouth recommendations."
    )
    # With MiniLM (or classical_strict=False), polish can take peer upgrades.
    refined = refine_sentence(source, min_wsd=0.10, polish=True)
    assert any(
        change.original.lower() == "encourage"
        and change.replacement.lower() == "promote"
        for change in refined.changes
    ) or "promote" in refined.text.lower()


def test_reputation_retention_stable_under_classical_strict():
    source = (
        "By consistently delivering high-quality service, organizations can "
        "build a strong reputation, increase customer retention, and "
        "encourage positive word-of-mouth recommendations."
    )
    plain = orchestrator.rewrite_document(
        source,
        lexical_polish=False,
        use_lexical_refinement=True,
        use_paraphrase=False,
        use_minilm_safety=False,
    )
    polished = orchestrator.rewrite_document(
        source,
        lexical_polish=True,
        use_lexical_refinement=True,
        use_paraphrase=False,
        use_minilm_safety=False,
    )
    for text in (plain.text.lower(), polished.text.lower()):
        assert "launch" not in text
        assert "functioning" not in text
        assert "found a strong reputation" not in text

def test_play_roles_is_not_swapped_to_act():
    resource = FakeWordnet(
        {
            ("play", "v"): [
                FakeSynset(
                    "fake-play-v",
                    "perform a role or function in a situation",
                    ["play", "act"],
                    examples=["they play important roles"],
                )
            ]
        }
    )
    source = (
        "Governments, businesses, and individuals all play important roles "
        "in preserving natural resources."
    )
    result = refine_sentence(source, min_wsd=0.10, wordnet=resource)
    assert "act important roles" not in result.text.lower()
    assert not any(change.original.lower() == "play" for change in result.changes)


def test_preserving_resources_is_not_swapped_to_continuing():
    source = (
        "Governments, businesses, and individuals all play important roles "
        "in preserving natural resources."
    )
    result = refine_sentence(source, min_wsd=0.14)
    assert "continuing natural resources" not in result.text.lower()
    assert "keeping natural resources" not in result.text.lower()
    low = result.text.lower()
    assert "preserving" in low or "conserve" in low or "protect" in low or "preserve" in low
    assert not any(
        change.replacement.lower().startswith(("continue", "keeping", "keep"))
        for change in result.changes
        if change.original.lower().startswith("preserv")
    )


def test_preserve_forests_is_not_swapped_to_continue():
    source = (
        "Communities should preserve forests, maintain clean rivers, and "
        "protect wildlife habitats."
    )
    result = refine_sentence(source, min_wsd=0.14)
    assert "continue forests" not in result.text.lower()
    assert "preserving" not in result.text.lower() or "preserve" in result.text.lower()
    assert not any(
        change.replacement.lower().startswith("continue") for change in result.changes
    )


def test_create_results_is_not_swapped_to_make():
    source = (
        "Effective communication is an essential skill that creates "
        "significant results in numerous workplaces."
    )
    result = refine_sentence(source, min_wsd=0.14)
    assert "makes" not in result.text.lower() or "creates" in result.text.lower()
    assert not any(
        change.original.lower() == "creates" and change.replacement.lower() == "makes"
        for change in result.changes
    )


def test_construct_can_simplify_to_build():
    source = "Workers construct bridges near the city every summer."
    result = refine_sentence(source, min_wsd=0.14)
    assert "build" in result.text.lower() or "construct" in result.text.lower()
    # Prefer build when WordNet offers the canonical construct/build sense.
    if result.changes:
        assert any(change.replacement.lower().startswith("build") for change in result.changes)


def test_assist_still_simplifies_to_help():
    source = "The report assists students with useful information."
    result = refine_sentence(source, min_wsd=0.14)
    assert "helps" in result.text.lower()
    assert any(change.replacement.lower() == "helps" for change in result.changes)


def test_purchase_still_simplifies_to_buy():
    source = "Companies should purchase reliable equipment carefully."
    result = refine_sentence(source, min_wsd=0.14)
    assert "buy" in result.text.lower() or "purchase" in result.text.lower()
    # Prefer buy when WordNet offers the clear purchase↔buy pair.
    if result.changes:
        assert any(change.replacement.lower() == "buy" for change in result.changes)


def test_sentence_initial_gerund_is_not_replaced():
    source = "Planning daily tasks allow people to complete their work more efficiently."
    result = refine_sentence(source, min_wsd=0.18)
    assert not any(change.original.lower() == "planning" for change in result.changes)
    assert "Projecting" not in result.text


def test_refinement_is_disabled_by_default(monkeypatch):
    monkeypatch.setattr(lexical, "_get_wordnet", provide_wordnet)
    result = orchestrator.rewrite_document(
        "The report assists students with useful information.",
        use_lexical_refinement=False,
        require_wording_change=False,
    )
    assert result.stats.lexical_refined == 0
    assert all(not record.lexical_changes for record in result.sentences)


def test_refinement_integrates_after_structural_rewrite(monkeypatch):
    monkeypatch.setattr(lexical, "_get_wordnet", provide_wordnet)
    result = orchestrator.rewrite_document(
        "The report assists students with useful information.",
        use_lexical_refinement=True,
        lexical_min_wsd=0.18,
        force_rewrite=False,
    )
    assert "helps" in result.text
    assert result.stats.lexical_refined == 1
    assert any(
        change.replacement == "helps"
        for change in result.sentences[0].lexical_changes
    )


def test_high_confidence_lexical_change_can_rescue_structural_skip(monkeypatch):
    monkeypatch.setattr(lexical, "_get_wordnet", provide_wordnet)
    source = (
        "The vital report that assists students with detailed information was "
        "reviewed by readers yesterday."
    )
    result = orchestrator.rewrite_document(
        source,
        use_lexical_refinement=True,
        lexical_min_wsd=0.18,
        force_rewrite=False,
    )
    assert "needed report" in result.text or "helps" in result.text
    assert result.sentences[0].status == "rewritten"
    assert "lexical" in (result.sentences[0].template_id or "") or result.stats.lexical_refined >= 1
    assert "needed report" in result.text or "helps" in result.text


def test_entity_loss_rolls_back_only_lexical_stage(monkeypatch):
    def unsafe_refinement(text: str, **_kwargs) -> LexicalResult:
        return LexicalResult(
            text=text.replace("Alice", "Someone"),
            changes=[
                LexicalChange(
                    original="Alice",
                    replacement="Someone",
                    token_index=0,
                    confidence=1.0,
                )
            ],
            confidence=1.0,
        )

    monkeypatch.setattr(orchestrator, "refine_sentence", unsafe_refinement)
    result = orchestrator.rewrite_document(
        "Alice visited Paris yesterday happily.",
        use_lexical_refinement=True,
    )
    assert "Alice" in result.text
    assert "Someone" not in result.text
    assert not result.sentences[0].lexical_changes


def test_special_blocks_remain_untouched_when_enabled(monkeypatch):
    monkeypatch.setattr(lexical, "_get_wordnet", provide_wordnet)
    table = "| Item | Detail |\n| --- | --- |\n| report | assists students |"
    result = orchestrator.rewrite_document(
        table,
        use_lexical_refinement=True,
    )
    assert result.text == table
    assert result.stats.lexical_refined == 0


def test_ten_thousand_word_lexical_batching_regression(monkeypatch):
    monkeypatch.setattr(lexical, "_get_wordnet", provide_wordnet)
    # Measure batching, not aggressive multi-pass wording throughput.
    monkeypatch.setattr(orchestrator, "ENGINE_CLASSICAL_AGGRESSIVE", False)
    sentence = "The report assists students with useful information today."
    paragraph = " ".join([sentence] * 50)
    source = "\n\n".join([paragraph] * 25)
    assert len(source.split()) >= 10_000

    result = orchestrator.rewrite_document(
        source,
        batch_paras=5,
        use_lexical_refinement=True,
        lexical_min_wsd=0.18,
        force_rewrite=False,
        use_paraphrase=False,
        use_minilm_safety=False,
    )
    assert result.input_words >= 10_000
    assert result.stats.batches == 5
    assert result.stats.lexical_refined > 0
    assert "helps" in result.text
    # CPU smoke budget: ~10k words / 1.2k sentences on a laptop can exceed 6 min.
    assert result.stats.seconds < 900