"""Source-preserving stress parsing and whole-word context windows."""
import re

VOWELS = 'аеёиоуыэюя'


WORD = re.compile('[А-Яа-яЁё]+')


def parse(text):
    plain, marks = [], set()
    for i, c in enumerate(text):
        if c == '+' and i+1 < len(text) and text[i+1].lower() in VOWELS:
            marks.add(len(plain))
        elif c == '\u0301' and plain and plain[-1].lower() in VOWELS:
            marks.add(len(plain)-1)
        else:
            plain.append(c)
    return ''.join(plain), marks


def encode(text, target, tokenizer, encoded=None):
    """Keep the complete original string and map each vowel to all BPE pieces."""
    if encoded is None:
        encoded = tokenizer(text, add_special_tokens=True, return_offsets_mapping=True,
                            return_attention_mask=False, truncation=False)
    offsets = encoded['offset_mapping']
    # Fast tokenizer byte pieces can share a character span: use both ends.
    covering = {}
    for token, (a, b) in enumerate(offsets):
        for p in range(a, b):
            if p in covering:
                covering[p][1] = token
            else:
                covering[p] = [token, token]
    words, candidates = [], []
    for wi, w in enumerate(WORD.finditer(text)):
        vowels = [p for p in range(w.start(), w.end()) if text[p].lower() in VOWELS]
        if not vowels:
            continue
        if any(p not in covering for p in range(w.start(), w.end())):
            raise ValueError('Tokenizer omitted characters inside a Russian word')
        first, last = covering[w.start()][0], covering[w.end()-1][1]
        off = len(candidates)
        for vi, pos in enumerate(vowels):
            a, b = covering[pos]
            candidates.append((a, b, pos, VOWELS.index(text[pos].lower()),
                               (pos-w.start())/len(w.group()), (w.end()-pos-1)/len(w.group())))
        label = int(target.get(wi, -100))
        if label != -100 and not 0 <= label < len(vowels):
            raise ValueError('Label outside candidate vowels')
        words.append((first, last, off, len(vowels), label, w.start(), w.end()))
    return encoded['input_ids'], words, candidates


def context_parts(text, tokenizer, limit=512):
    """Split only on whitespace, never truncate or cut a Russian word."""
    if len(tokenizer.encode(text, add_special_tokens=True, verbose=False)) <= limit:
        return [(0, len(text))]
    spans = [(m.start(), m.end()) for m in re.finditer(r'\S+', text)]
    parts = []
    start_word = 0
    while start_word < len(spans):
        lo, hi = start_word+1, len(spans)
        best = None
        while lo <= hi:
            mid = (lo+hi)//2
            a = 0 if start_word == 0 else spans[start_word][0]
            b = len(text) if mid == len(spans) else spans[mid][0]
            if len(tokenizer.encode(text[a:b], add_special_tokens=True, verbose=False)) <= limit:
                best = (a, b, mid); lo = mid+1
            else:
                hi = mid-1
        if best is None:
            raise ValueError('One unbroken word exceeds the encoder context')
        a, b, start_word = best
        parts.append((a, b))
    assert ''.join(text[a:b] for a, b in parts) == text
    return parts