rhyme.py

# rhyme.py
from __future__ import annotations

"""Utilities for rhyme extraction and classification."""

from dataclasses import dataclass
from typing import List, Optional, Set, Tuple

from syllables import CYRILLIC_VOWELS, normalize_text

_PUNCT_STRIP = " \t\r\n.,!?;:'\"()[]{}<>«»—–-…/\\|@#$%^&*_+=~`"


@dataclass(frozen=True)
class RhymeMark:
    """Classification result for a rhyme comparison between two lines."""

    kind: str  # "good" / "mid" / "bad"


def extract_last_word(line: str) -> str:
    """Return the last word in a line, stripped of punctuation."""
    text = (line or "").strip()
    if not text:
        return ""
    text = text.rstrip(_PUNCT_STRIP)
    if not text:
        return ""
    parts = text.split()
    if not parts:
        return ""
    return parts[-1].strip(_PUNCT_STRIP)


def normalize_word_for_rhyme(word: str) -> str:
    """Normalize a word for rhyme comparison."""
    normalized = normalize_text(word)
    while normalized.endswith(("ь", "ъ")):
        normalized = normalized[:-1]
    return normalized


def rhyme_tail(word: str, vowels: Optional[Set[str]] = None) -> str:
    """Return the tail of a word starting from its last vowel."""
    normalized = normalize_word_for_rhyme(word)
    if not normalized:
        return ""
    vowel_set = vowels or CYRILLIC_VOWELS
    last_vowel_index = -1
    for idx in range(len(normalized) - 1, -1, -1):
        if normalized[idx] in vowel_set:
            last_vowel_index = idx
            break
    if last_vowel_index < 0:
        return ""
    return normalized[last_vowel_index:]


def last_vowel(word: str, vowels: Optional[Set[str]] = None) -> str:
    """Return the last vowel found in the word."""
    normalized = normalize_word_for_rhyme(word)
    vowel_set = vowels or CYRILLIC_VOWELS
    for idx in range(len(normalized) - 1, -1, -1):
        if normalized[idx] in vowel_set:
            return normalized[idx]
    return ""


def rhyme_marks_for_block(lines: List[str]) -> List[RhymeMark]:
    """Classify rhymes for a block of lyric lines."""
    normalized_lines = [line.replace("|", "") for line in lines]
    if not normalized_lines:
        return []

    vowel_set = CYRILLIC_VOWELS
    marks: List[RhymeMark] = [RhymeMark("bad") for _ in normalized_lines]

    def compare(index_i: int, index_j: int) -> Tuple[RhymeMark, RhymeMark]:
        word_i = extract_last_word(normalized_lines[index_i])
        word_j = extract_last_word(normalized_lines[index_j])
        if not word_i or not word_j:
            return (RhymeMark("bad"), RhymeMark("bad"))

        tail_i = rhyme_tail(word_i, vowel_set)
        tail_j = rhyme_tail(word_j, vowel_set)
        if tail_i and tail_j and tail_i == tail_j:
            return (RhymeMark("good"), RhymeMark("good"))

        vowel_i = last_vowel(word_i, vowel_set)
        vowel_j = last_vowel(word_j, vowel_set)
        if vowel_i and vowel_j and vowel_i == vowel_j:
            return (RhymeMark("mid"), RhymeMark("mid"))

        return (RhymeMark("bad"), RhymeMark("bad"))

    for idx in range(0, len(normalized_lines), 2):
        if idx + 1 < len(normalized_lines):
            mark_i, mark_j = compare(idx, idx + 1)
            marks[idx] = mark_i
            marks[idx + 1] = mark_j
        else:
            marks[idx] = RhymeMark("bad")

    return marks


__all__ = [
    "RhymeMark",
    "extract_last_word",
    "normalize_word_for_rhyme",
    "rhyme_tail",
    "last_vowel",
    "rhyme_marks_for_block",
]