"""Collapse degenerate consecutive n-gram repetition loops in ASR hypotheses.

Extends the single-word repetition cap (max 3) to phrase loops: any n-gram
(n>=2) immediately repeating more than `max_rep` times is collapsed. Greedy
left-to-right per n, longest n first, deterministic — part of the model's
text post-processing, applied uniformly across datasets.
"""


def _collapse(words, n, max_rep):
    out = []
    i = 0
    while i < len(words):
        if len(out) >= n and i + n <= len(words) and words[i:i + n] == out[-n:]:
            reps = 1
            while len(out) >= n * (reps + 1) and out[-n * (reps + 1):len(out) - n * reps] == out[-n:]:
                reps += 1
            if reps >= max_rep:
                i += n
                continue
        out.append(words[i])
        i += 1
    return out


def fix_ngram_loops(text, max_n=16):
    """max_n=16 (was 8): batched decoding produced a period-10 loop
    ("and where the principles of the rule of law are respected" x ~90,
    VoxPopuli, 2026-06-12) that max_n=8 missed. The cap-sensitivity sweep
    shows max_n in {4,8,16} yields an identical macro on loop-free and
    short-period outputs, so widening only extends coverage."""
    words = text.split()
    for n in range(max_n, 0, -1):
        words = _collapse(words, n, 3 if n == 1 else 2)
    return " ".join(words)


if __name__ == "__main__":
    assert fix_ngram_loops("in the " * 190 + "end") == "in the in the end"
    assert fix_ngram_loops("he he he he he") == "he he he"
    assert fix_ngram_loops("thank you thank you") == "thank you thank you"
    assert fix_ngram_loops("a b c a b c a b c a b c") == "a b c a b c"
    assert fix_ngram_loops("normal sentence with no loops at all") == "normal sentence with no loops at all"
    assert fix_ngram_loops("very very good") == "very very good"
    ten = "and where the principles of the rule of law are respected"
    looped = "intro words here " + " ".join([ten] * 90)
    assert fix_ngram_loops(looped) == "intro words here " + " ".join([ten] * 2)
    print("all tests pass")