tests/unit/core/test_normalization.py

# License: Apache-2.0 License
# Created by: Patrick Lumbantobing, VertoX-AI
# Copyright (c) 2026 VertoX-AI. All rights reserved.
#
# This work is licensed under the Apache-2.0 License.
# To view a copy of this license, visit
# https://www.apache.org/licenses/LICENSE-2.0

from __future__ import annotations

import math

import pytest

from speech_text_data_aligners.core.errors import (
    TranscriptError,
    TranscriptMappingError,
)
from speech_text_data_aligners.core.models import LexicalWord
from speech_text_data_aligners.core.normalization import (
    ObservedWordUnit,
    TextPlan,
    TextPlanUnit,
    comparison_key,
    lexicalize,
    make_text_plan,
    reconcile_word_units,
)


def test_lexicalize_preserves_punctuation_sensitive_surfaces_and_spans() -> None:
    transcript = "“Hello,” U.S.A. can’t re—enter 3.14 café!"
    words = lexicalize(transcript)
    assert [word.text for word in words] == [
        "Hello",
        "U.S.A.",
        "can’t",
        "re—enter",
        "3.14",
        "café",
    ]
    assert all(
        transcript[word.start_char : word.end_char] == word.text for word in words
    )
    assert [word.index for word in words] == list(range(6))


@pytest.mark.parametrize(
    ("transcript", "expected"),
    [
        ("Value 3.14.", "3.14"),
        ("Version 1.2.3.", "1.2.3"),
        ("Visit example.com.", "example.com"),
        ("Country U.S.", "U.S."),
        ('Quoted "3.14."', "3.14"),
        ("Bracketed (example.com.)", "example.com"),
        ('Abbreviation "U.S."', "U.S."),
    ],
)
def test_lexicalize_distinguishes_terminal_punctuation_from_abbreviation_dots(
    transcript: str, expected: str
) -> None:
    word = lexicalize(transcript)[-1]
    assert word.text == expected
    assert transcript[word.start_char : word.end_char] == expected


@pytest.mark.parametrize("transcript", ["", "   ", "... — !!!"])
def test_lexicalize_rejects_empty_or_nonlexical_text(transcript: str) -> None:
    with pytest.raises(TranscriptError):
        lexicalize(transcript)


def test_comparison_key_documents_case_unicode_and_punctuation_folding() -> None:
    assert comparison_key("Ｕ.S.—Can’T") == "uscan‘t".replace("‘", "")
    assert comparison_key("Straße") == "strasse"
    assert comparison_key("---") == ""


def test_text_plan_supports_expansions_aliases_and_repeated_words() -> None:
    plan = make_text_plan(
        "Go 12 go",
        expansions={1: ("twelve",)},
        aliases={1: ("dozen",)},
    )
    assert plan.emitted_text == "Go twelve go"
    assert [unit.lexical_index for unit in plan.units] == [0, 1, 2]
    observed = (
        ObservedWordUnit("go", 0.0, 0.1, 0.8),
        ObservedWordUnit("dozen", 0.2, 0.4, 0.6),
        ObservedWordUnit("GO", 0.5, 0.7, 0.9),
    )
    result = reconcile_word_units(plan, observed)
    assert [item.lexical_word.text for item in result] == ["Go", "12", "go"]
    assert [item.start_s for item in result] == [0.0, 0.2, 0.5]


def test_reconcile_collapses_multiple_units_and_weights_confidence() -> None:
    plan = make_text_plan("21", expansions={0: ("twenty", "one")})
    result = reconcile_word_units(
        plan,
        (
            ObservedWordUnit("twenty", 0.1, 0.3, 0.5),
            ObservedWordUnit("one", 0.3, 0.7, 1.0),
        ),
    )
    assert result[0].start_s == 0.1
    assert result[0].end_s == 0.7
    assert result[0].confidence == pytest.approx((0.5 * 0.2 + 1.0 * 0.4) / 0.6)
    assert result[0].attributes == (("unit_0", "twenty"), ("unit_1", "one"))


def test_reconcile_rejects_unexpected_observations_by_default() -> None:
    repeated_plan = make_text_plan("go go")
    repeated_observed = (
        ObservedWordUnit("go", 0.0, 0.1),
        ObservedWordUnit("go", 0.2, 0.3),
        ObservedWordUnit("go", 0.4, 0.5),
    )
    expanded_plan = make_text_plan("Start 21", expansions={1: ("twenty", "one")})
    expanded_observed = (
        ObservedWordUnit("noise", 0.0, 0.05),
        ObservedWordUnit("start", 0.1, 0.2),
        ObservedWordUnit("unmatched", 0.25, 0.3),
        ObservedWordUnit("twenty", 0.4, 0.6),
        ObservedWordUnit("one", 0.6, 0.9),
        ObservedWordUnit("tail", 1.0, 1.1),
    )

    for plan, observed in (
        (repeated_plan, repeated_observed),
        (expanded_plan, expanded_observed),
    ):
        with pytest.raises(TranscriptMappingError, match="cannot cover"):
            reconcile_word_units(plan, observed)


@pytest.mark.parametrize(
    ("observed", "ignorable_index", "expected_starts"),
    [
        (
            (
                ObservedWordUnit("noise", 0.0, 0.05),
                ObservedWordUnit("start", 0.1, 0.2),
                ObservedWordUnit("end", 0.3, 0.4),
            ),
            0,
            [0.1, 0.3],
        ),
        (
            (
                ObservedWordUnit("start", 0.0, 0.1),
                ObservedWordUnit("noise", 0.15, 0.2),
                ObservedWordUnit("end", 0.3, 0.4),
            ),
            1,
            [0.0, 0.3],
        ),
        (
            (
                ObservedWordUnit("start", 0.0, 0.1),
                ObservedWordUnit("end", 0.2, 0.3),
                ObservedWordUnit("noise", 0.4, 0.5),
            ),
            2,
            [0.0, 0.2],
        ),
    ],
    ids=["leading", "middle", "trailing"],
)
def test_reconcile_dp_allows_one_explicitly_authorized_skip(
    observed: tuple[ObservedWordUnit, ...],
    ignorable_index: int,
    expected_starts: list[float],
) -> None:
    result = reconcile_word_units(
        make_text_plan("start end"),
        observed,
        ignorable_observed_indices=frozenset({ignorable_index}),
    )
    assert [word.start_s for word in result] == expected_starts


def test_reconcile_dp_allows_multiple_authorized_skips_and_expansion() -> None:
    plan = make_text_plan("Start 21", expansions={1: ("twenty", "one")})
    result = reconcile_word_units(
        plan,
        (
            ObservedWordUnit("noise", 0.0, 0.05),
            ObservedWordUnit("start", 0.1, 0.2),
            ObservedWordUnit("unmatched", 0.25, 0.3),
            ObservedWordUnit("twenty", 0.4, 0.6),
            ObservedWordUnit("one", 0.6, 0.9),
            ObservedWordUnit("tail", 1.0, 1.1),
        ),
        ignorable_observed_indices=frozenset({0, 2, 5}),
    )
    assert [(word.start_s, word.end_s) for word in result] == [
        (0.1, 0.2),
        (0.4, 0.9),
    ]
    assert result[1].attributes == (("unit_0", "twenty"), ("unit_1", "one"))


def test_reconcile_same_spelling_extra_requires_index_authorization() -> None:
    plan = make_text_plan("go go")
    observed = (
        ObservedWordUnit("go", 0.0, 0.1),
        ObservedWordUnit("GO", 0.2, 0.3),
        ObservedWordUnit("go", 0.4, 0.5),
    )

    with pytest.raises(TranscriptMappingError, match="cannot cover"):
        reconcile_word_units(plan, observed)

    result = reconcile_word_units(
        plan, observed, ignorable_observed_indices=frozenset({2})
    )
    assert [word.start_s for word in result] == [0.0, 0.2]


def test_reconcile_dp_uses_stable_earliest_authorized_path() -> None:
    plan = make_text_plan("go go")
    observed = tuple(
        ObservedWordUnit("go", index / 10, index / 10 + 0.05) for index in range(4)
    )
    results = tuple(
        reconcile_word_units(
            plan,
            observed,
            ignorable_observed_indices=frozenset({0, 2, 3}),
        )
        for _ in range(5)
    )
    assert all(result == results[0] for result in results)
    assert [word.start_s for word in results[0]] == [0.0, 0.1]


def test_reconcile_dp_matches_aliases_with_an_authorized_insertion() -> None:
    plan = make_text_plan("a 12 b", aliases={1: ("dozen",)})
    result = reconcile_word_units(
        plan,
        (
            ObservedWordUnit("a", 0.0, 0.1),
            ObservedWordUnit("extra", 0.1, 0.2),
            ObservedWordUnit("dozen", 0.2, 0.3),
            ObservedWordUnit("b", 0.3, 0.4),
        ),
        ignorable_observed_indices=frozenset({1}),
    )
    assert [word.lexical_word.text for word in result] == ["a", "12", "b"]


def test_reconcile_rejects_unauthorized_extra_when_another_skip_is_authorized() -> None:
    with pytest.raises(TranscriptMappingError, match="cannot cover"):
        reconcile_word_units(
            make_text_plan("a b"),
            (
                ObservedWordUnit("authorized", 0.0, 0.05),
                ObservedWordUnit("a", 0.1, 0.2),
                ObservedWordUnit("unauthorized", 0.25, 0.3),
                ObservedWordUnit("b", 0.4, 0.5),
            ),
            ignorable_observed_indices=frozenset({0}),
        )


@pytest.mark.parametrize("invalid_index", [-1, 1])
def test_reconcile_rejects_out_of_range_skip_authorization(
    invalid_index: int,
) -> None:
    with pytest.raises(TranscriptMappingError, match="out of range"):
        reconcile_word_units(
            make_text_plan("word"),
            (ObservedWordUnit("word", 0.0, 0.1),),
            ignorable_observed_indices=frozenset({invalid_index}),
        )


@pytest.mark.parametrize("invalid_index", ["0", 0.0, True])
def test_reconcile_rejects_non_integer_skip_authorization(
    invalid_index: object,
) -> None:
    with pytest.raises(TranscriptMappingError, match="only integers"):
        reconcile_word_units(
            make_text_plan("word"),
            (ObservedWordUnit("word", 0.0, 0.1),),
            ignorable_observed_indices=frozenset({invalid_index}),  # type: ignore[arg-type]
        )


def test_reconcile_dp_rejects_missing_required_planned_unit() -> None:
    with pytest.raises(TranscriptMappingError, match="cover"):
        reconcile_word_units(
            make_text_plan("required missing"),
            (
                ObservedWordUnit("required", 0.0, 0.1),
                ObservedWordUnit("unmatched", 0.2, 0.3),
            ),
        )


def test_reconcile_zero_duration_confidences_and_missing_confidence() -> None:
    plan = make_text_plan("7", expansions={0: ("seven", "seven")})
    result = reconcile_word_units(
        plan,
        (
            ObservedWordUnit("seven", 0.2, 0.2, 0.2),
            ObservedWordUnit("seven", 0.2, 0.2, 0.6),
        ),
    )
    assert result[0].confidence == pytest.approx(0.4)
    no_confidence = reconcile_word_units(
        make_text_plan("word"), (ObservedWordUnit("word", 0.0, 0.1),)
    )
    assert no_confidence[0].confidence is None


@pytest.mark.parametrize(
    "observed",
    [
        (),
        (ObservedWordUnit("wrong", 0.0, 0.1),),
        (ObservedWordUnit("word", math.nan, 0.1),),
        (ObservedWordUnit("word", 0.0, 0.1, math.inf),),
        (
            ObservedWordUnit("ignored", math.inf, 0.1),
            ObservedWordUnit("word", 0.2, 0.3),
        ),
    ],
)
def test_reconcile_rejects_undercoverage_spelling_and_nonfinite_values(
    observed: tuple[ObservedWordUnit, ...],
) -> None:
    with pytest.raises(TranscriptMappingError):
        reconcile_word_units(make_text_plan("word"), observed)


def test_reconcile_validates_an_authorized_skipped_observation() -> None:
    with pytest.raises(TranscriptMappingError, match="non-finite timestamp"):
        reconcile_word_units(
            make_text_plan("word"),
            (
                ObservedWordUnit("native-artifact", math.inf, math.inf),
                ObservedWordUnit("word", 0.0, 0.1),
            ),
            ignorable_observed_indices=frozenset({0}),
        )


def test_reconcile_rejects_uncovered_lexical_word() -> None:
    plan = TextPlan(
        "a b",
        (LexicalWord(0, "a", 0, 1), LexicalWord(1, "b", 2, 3)),
        "a",
        (TextPlanUnit("a", 0),),
    )
    with pytest.raises(TranscriptMappingError, match="no timed"):
        reconcile_word_units(plan, (ObservedWordUnit("a", 0.0, 0.1),))


@pytest.mark.parametrize(
    "plan",
    [
        TextPlan("a", (LexicalWord(3, "a", 0, 1),), "a", (TextPlanUnit("a", 0),)),
        TextPlan("a", (LexicalWord(0, "a", 0, 1),), "a", (TextPlanUnit("a", 2),)),
        TextPlan(
            "a b",
            (LexicalWord(0, "a", 0, 1), LexicalWord(1, "b", 2, 3)),
            "b a",
            (TextPlanUnit("b", 1), TextPlanUnit("a", 0)),
        ),
    ],
)
def test_reconcile_rejects_invalid_plan_ownership(plan: TextPlan) -> None:
    with pytest.raises(TranscriptMappingError, match="text plan"):
        reconcile_word_units(plan, (ObservedWordUnit("a", 0.0, 0.1),))


@pytest.mark.parametrize(
    ("expansions", "aliases"),
    [({3: ("bad",)}, None), ({0: ()}, None), ({0: ("",)}, None), (None, {2: ("x",)})],
)
def test_make_text_plan_rejects_invalid_mapping_configuration(
    expansions: dict[int, tuple[str, ...]] | None,
    aliases: dict[int, tuple[str, ...]] | None,
) -> None:
    with pytest.raises(TranscriptError):
        make_text_plan("one", expansions=expansions, aliases=aliases)