src/inference.py

#!/usr/bin/env python3
"""
Inference script for MHG text normalization.
Apply trained model to new, unannotated MHG texts.

Supports sliding-window stride inference for long texts that exceed
the model's maximum sequence length. When stride is enabled, long
inputs are split into overlapping windows, each normalized independently,
and predictions are merged preferring window centers (where context is
richest).
"""

import os
import yaml
import torch
import argparse
from pathlib import Path
from typing import List, Optional, Tuple
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import re
import difflib
from tqdm import tqdm

# Editorial punctuation characters to strip before normalization and restore after.
# These appear in edited/printed text but were NOT present in the training data.
# MHG abbreviation markers (hyphens '-', apostrophes '\u2019') are NOT included
# here because they are meaningful for the normalization task.
_PUNCT_CHARS = frozenset(".,;:!?()[]«»‹›\u201c\u201d\u201e\"\u201a\u2018\u2019\u2014\u2013<>")


def _attach_en_proclitic(text: str) -> str:
    """Attach free-standing negation proclitic 'en' to the following word.

    In Middle High German, the negation proclitic 'en' appears as a separate
    token before the word it negates. This post-processing step joins them
    directly without a separator, e.g. 'en guot' → 'enguot'.

    Capitalization of the proclitic is preserved: 'En guot' → 'Enguot'.

    Only attaches when 'en' appears as a standalone, case-insensitive word
    (bounded by whitespace or start/end of string) and is followed by another
    word. Punctuation after 'en' prevents attachment.

    Args:
        text: Input text to process

    Returns:
        Text with 'en' proclitics attached to the following word
    """
    # Match standalone "en" (case-insensitive) as a whole word,
    # followed by whitespace and a non-whitespace character.
    pattern = re.compile(r'(?<!\S)(en)(?!\S)\s+(\S)', re.IGNORECASE)

    def _replace(m):
        en_word = m.group(1)
        next_char = m.group(2)
        return f"{en_word}{next_char}"

    return pattern.sub(_replace, text)


def _lenition_t_after_n_before_vowel(text: str) -> str:
    """MHG Lenition: /t/ → /d/ after /n/ and before a vowel.

    In Middle High German, the voiceless stop /t/ becomes voiced /d/ when
    it appears between /n/ and a vowel (within the same word). This covers
    all MHG vowels including those with diacritics (ä, ö, ü, ê, î, etc.)
    and base vowels followed by combining diacritical marks.

    Capitalization is preserved: 'Ntô' → 'Ndô', 'nTô' → 'ndô'.

    Args:
        text: Input text to process

    Returns:
        Text with lenited t → d where applicable
    """
    # Vowel set: MHG base vowels plus all common diacritic variants
    vowel_chars = 'aeiouyAEIOUY'
    vowel_chars += 'äöüÄÖÜ'           # umlaut
    vowel_chars += 'âêîôûÂÊÎÔÛ'      # circumflex
    vowel_chars += 'āēīōūĀĒĪŌŪ'      # macron
    vowel_chars += 'åæœøÅÆŒØ'        # other
    # Match 'n' + 't' followed by a vowel (optionally with combining diacritics)
    pattern = re.compile(
        r'([nN])([tT])(?=[' + re.escape(vowel_chars) + r'][\u0300-\u036f]*)'
    )

    def _replace(m):
        n_char = m.group(1)
        t_char = m.group(2)
        return n_char + ('d' if t_char == 't' else 'D')

    return pattern.sub(_replace, text)


def _lenition_t_after_l_before_vowel(text: str) -> str:
    """MHG Lenition: /t/ → /d/ after /l/ and before a vowel.

    In Middle High German, the voiceless stop /t/ becomes voiced /d/ when
    it appears between /l/ and a vowel (within the same word). This covers
    all MHG vowels including those with diacritics (ä, ö, ü, ê, î, etc.)
    and base vowels followed by combining diacritical marks.

    Capitalization is preserved: 'Ltô' → 'Ldô', 'lTô' → 'ldô'.

    Args:
        text: Input text to process

    Returns:
        Text with lenited t → d where applicable
    """
    # Vowel set: MHG base vowels plus all common diacritic variants
    vowel_chars = 'aeiouyAEIOUY'
    vowel_chars += 'äöüÄÖÜ'           # umlaut
    vowel_chars += 'âêîôûÂÊÎÔÛ'      # circumflex
    vowel_chars += 'āēīōūĀĒĪŌŪ'      # macron
    vowel_chars += 'åæœøÅÆŒØ'        # other
    # Match 'l' + 't' followed by a vowel (optionally with combining diacritics)
    pattern = re.compile(
        r'([lL])([tT])(?=[' + re.escape(vowel_chars) + r'][\u0300-\u036f]*)'
    )

    def _replace(m):
        l_char = m.group(1)
        t_char = m.group(2)
        return l_char + ('d' if t_char == 't' else 'D')

    return pattern.sub(_replace, text)


def _niet_to_niht(text: str) -> str:
    """Change freestanding 'niet'/'niut' to 'niht'.

    In some MHG texts, the negation word 'niet' (or variant 'niut') appears
    as a freestanding token but the preferred normalized form is 'niht'.
    This post-processing step replaces standalone occurrences of 'niet' and
    'niut' (case-insensitive) with 'niht', preserving the original
    capitalization pattern.

    Handles words immediately followed by editorial punctuation (from
    _PUNCT_CHARS) by capturing and re-attaching the trailing punctuation.

    Examples:
        'niet' → 'niht'
        'niut' → 'niht'
        'Niet' → 'Niht'
        'Niut' → 'Niht'
        'NIET' → 'NIHT'
        'NIUT' → 'NIHT'
        'niet.' → 'niht.'
        'niet,' → 'niht,'
        'niet.)' → 'niht.)'

    Args:
        text: Input text to process

    Returns:
        Text with freestanding 'niet'/'niut' changed to 'niht'
    """
    # Build character class of all editorial punctuation characters so we can
    # explicitly match trailing punctuation and re-attach it after replacement.
    # This is necessary because post-processing runs AFTER punctuation restoration,
    # so words like "niet." have punctuation attached.
    punct_class = re.escape(''.join(_PUNCT_CHARS))
    # Match standalone niet/niut (bounded by non-word on left) followed by
    # zero or more punctuation characters and then a non-word boundary.
    pattern = re.compile(
        r'(?<!\S)(ni[eu]t)([' + punct_class + r']*)(?!\S)', re.IGNORECASE
    )

    def _replace(m):
        word = m.group(1)
        trailing = m.group(2)
        if word in ('niet', 'niut'):
            result = 'niht'
        elif word in ('Niet', 'Niut'):
            result = 'Niht'
        elif word in ('NIET', 'NIUT'):
            result = 'NIHT'
        else:
            # Mixed case — preserve first letter case
            result = (word[0].upper() if word[0].isupper() else word[0].lower()) + 'iht'
        return result + trailing

    return pattern.sub(_replace, text)


def _common_apocopes(text: str) -> str:
    """Apply common MHG apocope (final vowel/consonant loss) rules.

    In Middle High German, certain words lose their final syllables in
    common usage. This post-processing step handles the most frequent
    apocopes:

    Examples:
        'wile' → 'wil'
        'vore' → 'vor'

    Args:
        text: Input text to process

    Returns:
        Text with common apocopes applied
    """
    # Build character class of all editorial punctuation so we can
    # capture trailing punctuation and re-attach it after replacement.
    punct_class = re.escape(''.join(_PUNCT_CHARS))

    replacements = {
        'vile': 'vil',
        'vore': 'vor',
        'wile': 'wil',
        'wole': 'wol',
    }

    # Build a single regex that matches any of the keys as whole words
    # (bounded by whitespace or start/end), followed by optional punctuation.
    pattern_str = r'(?<!\S)(' + '|'.join(re.escape(k) for k in replacements.keys()) + r')([' + punct_class + r']*)(?!\S)'
    pattern = re.compile(pattern_str, re.IGNORECASE)

    def _replace(m):
        word = m.group(1)
        trailing = m.group(2)
        result = replacements.get(word.lower(), word)
        # Preserve capitalization pattern from the original word
        if word[0].isupper() and result[0].islower():
            result = result[0].upper() + result[1:]
        return result + trailing

    return pattern.sub(_replace, text)


def _capitalize_first_alpha(word: str) -> str:
    """Capitalize the first alphabetic character in a word.
    
    Handles words that may start with non-alphabetic characters
    (e.g., MHG abbreviation markers like hyphens or apostrophes,
    or leading punctuation that was not stripped).
    
    Args:
        word: Word to capitalize
    
    Returns:
        Word with first alphabetic character uppercased
    """
    for i, ch in enumerate(word):
        if ch.isalpha():
            return word[:i] + ch.upper() + word[i+1:]
    return word  # No alphabetic characters found; return unchanged


def _select_device() -> torch.device:
    """Select the best available device: CUDA > ROCm > DirectML > MPS > CPU."""
    if torch.cuda.is_available():
        return torch.device('cuda')
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return torch.device('mps')
    try:
        import torch_directml
        device = torch_directml.device()
        return torch.device(device)
    except ImportError:
        pass
    return torch.device('cpu')


class MHGNormalizer:
    """Normalizer for Middle High German text."""
    
    def __init__(self, model_path: str, config_path: Optional[str] = None, defer_gpu: bool = False):
        """
        Initialize normalizer with trained model.
        
        Args:
            model_path: Path to trained model checkpoint
            config_path: Path to configuration file (optional)
        """
        self.device = _select_device()
        print(f"Using device: {self.device}")
        
        # Load configuration if provided
        if config_path:
            with open(config_path, 'r', encoding='utf-8') as f:
                self.config = yaml.safe_load(f)
        else:
            # Default configuration
            self.config = {
                'inference': {
                    'max_length': 512,
                    'num_beams': 5,
                    'length_penalty': 0.6,
                    'do_sample': False,
                    'top_p': 0.9,
                    'temperature': 1.0,
                    'stride': True,
                    'stride_window_tokens': 400,
                    'stride_step_tokens': 300,
                }
            }
        
        # Ensure stride config exists (for configs created before stride was added)
        inference_config = self.config.setdefault('inference', {})
        if 'stride' not in inference_config:
            inference_config['stride'] = True
        if 'stride_window_tokens' not in inference_config:
            inference_config['stride_window_tokens'] = 400
        if 'stride_step_tokens' not in inference_config:
            inference_config['stride_step_tokens'] = 300
        
        # Ensure preserve_punctuation and preserve_capitalization defaults
        if 'preserve_punctuation' not in inference_config:
            inference_config['preserve_punctuation'] = True
        if 'preserve_capitalization' not in inference_config:
            inference_config['preserve_capitalization'] = True

        # Ensure post_processing defaults
        post_config = self.config.setdefault('post_processing', {})
        if 'attach_en_proclitic' not in post_config:
            post_config['attach_en_proclitic'] = False
        if 'lenition_t_after_n' not in post_config:
            post_config['lenition_t_after_n'] = False
        if 'lenition_t_after_l' not in post_config:
            post_config['lenition_t_after_l'] = False
        if 'niet_to_niht' not in post_config:
            post_config['niet_to_niht'] = False
        if 'common_apocopes' not in post_config:
            post_config['common_apocopes'] = False
        
        # Load model and tokenizer
        print(f"\nLoading model from: {model_path}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)

        # Detect dtype from model config (FP16 models have torch_dtype=float16)
        import json as _json
        from pathlib import Path as _Path
        _cfg = _Path(model_path) / "config.json"
        if _cfg.exists():
            _cfg_data = _json.loads(_cfg.read_text())
            _model_dtype = torch.float16 if _cfg_data.get("torch_dtype") == "float16" else None
        else:
            _model_dtype = None

        # CPU lacks efficient FP16 support — load as FP32 there. On CUDA/ZeroGPU
        # load FP16 weights as bfloat16, which is numerically stable for T5/mT5
        # (FP16 tends to overflow/NaN in T5) and fast on modern GPUs.
        # When defer_gpu=True (ZeroGPU), the model is preloaded to CPU at
        # startup and moved to GPU later inside @spaces.GPU — load as BF16
        # so the GPU transfer is fast and inference runs in BF16.
        if _model_dtype == torch.float16:
            if str(self.device) == "cpu" and not defer_gpu:
                print("  CPU detected — loading FP16 model as FP32 (CPU has no FP16 acceleration)")
                _load_dtype = torch.float32
            else:
                _where = "GPU" if str(self.device) != "cpu" else "CPU (GPU deferred)"
                print(f"  Loading FP16 model as bfloat16 ({_where})")
                _load_dtype = torch.bfloat16
        else:
            _load_dtype = _model_dtype

        # Do NOT pass tie_word_embeddings=False — T5/mT5 always forces
        # tie_word_embeddings=True in T5Config.__post_init__. Passing False
        # prevents HuggingFace from tying encoder/decoder embed_tokens to
        # shared.weight when loading checkpoints, causing broken inference.
        self.model = AutoModelForSeq2SeqLM.from_pretrained(
            model_path,
            torch_dtype=_load_dtype,
            low_cpu_mem_usage=True,
        )
        self.model.to(self.device)
        self.model.eval()
        
        print("Model loaded successfully!\n")

    def _ensure_on_device(self):
        """Move the model to the best available device if it has changed.

        On ZeroGPU, __init__ runs at startup when no GPU is visible, so the
        model loads to CPU. This is called inside @spaces.GPU (where CUDA is
        available) to transfer the model to the GPU. Idempotent: if the model
        is already on the target device, nothing happens.
        """
        target = _select_device()
        if target != self.device:
            print(f"Moving model from {self.device} to {target} …")
            self.device = target
            self.model.to(self.device)
    
    @staticmethod
    def _extract_punctuation(text: str):
        """Extract editorial punctuation from text for preservation during normalization.
        
        Strips leading and trailing editorial punctuation characters from each word,
        returning the stripped text and a punctuation map for later restoration.
        
        Words that consist entirely of punctuation characters are absorbed into
        the trailing punctuation of the previous word (or the leading punctuation
        of the next word if they appear at the start of the text).
        
        Args:
            text: Input text that may contain editorial punctuation
        
        Returns:
            tuple: (stripped_text, punct_map) where:
                - stripped_text: text with editorial punctuation removed from word edges
                - punct_map: list of (leading_punct, trailing_punct) tuples, one per
                  word in stripped_text
        """
        words = text.split()
        punct_map = []
        stripped_words = []
        pending_leading = ''
        
        for word in words:
            # Extract leading punctuation
            leading = ''
            i = 0
            while i < len(word) and word[i] in _PUNCT_CHARS:
                leading += word[i]
                i += 1
            
            # Extract trailing punctuation
            trailing = ''
            j = len(word) - 1
            while j >= i and word[j] in _PUNCT_CHARS:
                trailing = word[j] + trailing
                j -= 1
            
            core = word[i:j+1]
            
            if core:
                # Prepend any pending leading punctuation (from standalone punct at start)
                leading = pending_leading + leading
                pending_leading = ''
                punct_map.append((leading, trailing))
                stripped_words.append(core)
            else:
                # Word is entirely punctuation (e.g., "..." or "—")
                # Absorb into adjacent word's punctuation. Since ' '.join() provides
                # only one space between words, we include an explicit space to
                # preserve the whitespace around standalone punctuation marks.
                if punct_map:
                    # Append to previous word's trailing punctuation with leading space
                    prev_leading, prev_trailing = punct_map[-1]
                    punct_map[-1] = (prev_leading, prev_trailing + ' ' + leading + trailing)
                else:
                    # At start of text — save as leading punctuation for next word
                    pending_leading += leading + trailing
        
        # If text ended with standalone punctuation, attach to last word's trailing
        if pending_leading and punct_map:
            prev_leading, prev_trailing = punct_map[-1]
            punct_map[-1] = (prev_leading, prev_trailing + pending_leading)
        
        stripped_text = ' '.join(stripped_words)
        return stripped_text, punct_map
    
    @staticmethod
    def _restore_punctuation(text: str, punct_map: list) -> str:
        """Restore editorial punctuation to normalized text.
        
        Uses position-based mapping: word i gets punct_map[i]. When the
        model produces a different number of words than the original,
        trailing punctuation from unmatched positions is appended to the
        last word.
        
        For across-lines mode where word counts may change significantly,
        use _restore_punctuation_aligned() instead, which uses character-level
        alignment to correctly map punctuation even when words split or merge.
        
        Args:
            text: Normalized text (without editorial punctuation)
            punct_map: List of (leading_punct, trailing_punct) tuples from
                       _extract_punctuation
        
        Returns:
            Text with editorial punctuation restored
        """
        words = text.split()
        result = []
        
        for i, word in enumerate(words):
            if i < len(punct_map):
                leading, trailing = punct_map[i]
                result.append(leading + word + trailing)
            else:
                # More words in output than in punct_map; no punctuation info
                result.append(word)
        
        # If model produced fewer words than original, collect trailing punctuation
        # from unmatched positions and append to the last output word.
        # This preserves sentence-final punctuation (e.g., trailing periods).
        if len(words) > 0 and len(words) < len(punct_map):
            extra_trailing = ''.join(t for _, t in punct_map[len(words):])
            if extra_trailing:
                result[-1] = result[-1] + extra_trailing
        
        return ' '.join(result)
    
    @staticmethod
    def _restore_punctuation_aligned(original_text: str, normalized_text: str,
                                      punct_map: list) -> str:
        """Restore editorial punctuation using character-level alignment.
        
        When the model changes word counts (e.g., "inalten" → "in alten"),
        position-based _restore_punctuation() fails because word indices shift.
        This method uses character-level sequence alignment to correctly map
        each normalized word to its corresponding original word, then applies
        the original word's punctuation.
        
        For example, if "inalten" (with no punctuation) splits into "in alten",
        character alignment maps "in" to the start of "inalten" and "alten" to
        the end, so both correctly inherit the punctuation of "inalten".
        
        Args:
            original_text: Original text with punctuation already stripped
            normalized_text: Normalized text (without editorial punctuation)
            punct_map: List of (leading_punct, trailing_punct) tuples,
                       one per word in original_text
        
        Returns:
            Text with editorial punctuation restored
        """
        normalized_words = normalized_text.split()
        if not normalized_words:
            return normalized_text
        
        original_words = original_text.split()
        if not original_words:
            return normalized_text
        
        # If word counts match, use simple position-based restoration
        if len(normalized_words) == len(original_words):
            return MHGNormalizer._restore_punctuation(normalized_text, punct_map)
        
        # Use character-level alignment to map each normalized word to its
        # corresponding original word index. This correctly handles word
        # splits like "inalten" → "in alten" because character alignment
        # maps "in" to the "in" prefix of "inalten".
        norm_to_orig_word = MHGNormalizer._align_normalized_to_original_charlevel(
            original_text, normalized_text
        )
        
        # Apply punctuation, handling word splits:
        # When multiple normalized words map to the same original word,
        # only the first gets leading punct and only the last gets trailing punct.
        result = []
        for i, word in enumerate(normalized_words):
            orig_idx = norm_to_orig_word[i]
            if orig_idx < len(punct_map):
                leading, trailing = punct_map[orig_idx]
            else:
                leading, trailing = '', ''
            
            # Check if this is the first/last normalized word mapping
            # to this original word (handles word splits)
            is_first = (i == 0 or norm_to_orig_word[i - 1] != orig_idx)
            is_last = (i == len(normalized_words) - 1 or
                       norm_to_orig_word[i + 1] != orig_idx)
            
            applied_leading = leading if is_first else ''
            applied_trailing = trailing if is_last else ''
            result.append(applied_leading + word + applied_trailing)
        
        # Collect trailing punctuation from unmatched original words
        # (original words that were deleted by the model)
        matched_orig = set(norm_to_orig_word)
        extra_trailing = ''
        for orig_idx in range(len(punct_map)):
            if orig_idx not in matched_orig:
                _, trailing = punct_map[orig_idx]
                extra_trailing += trailing
        if extra_trailing and result:
            result[-1] = result[-1] + extra_trailing
        
        return ' '.join(result)
    
    @staticmethod
    def _extract_capitalization(text: str):
        """Extract capitalization pattern from text for preservation during normalization.
        
        Records which words start with an uppercase letter, then lowercases
        the text so the model receives consistent lowercase input.
        
        Words whose first alphabetic character is uppercase are marked in the
        cap_map. Non-alphabetic leading characters (e.g., punctuation that
        wasn't stripped, MHG abbreviation markers) are skipped when finding
        the first alphabetic character.
        
        This method should be called AFTER _extract_punctuation so that
        the cap_map indices align with the punct_map indices (both indexed
        by words in the punctuation-stripped text).
        
        Args:
            text: Input text (typically with punctuation already stripped)
        
        Returns:
            tuple: (lowercased_text, cap_map) where:
                - lowercased_text: the text converted to lowercase
                - cap_map: list of bools, one per word; True if the word's
                  first alphabetic character was uppercase
        """
        words = text.split()
        cap_map = []
        for word in words:
            has_cap = False
            for ch in word:
                if ch.isalpha():
                    has_cap = ch.isupper()
                    break
            cap_map.append(has_cap)
        lowered_text = text.lower()
        return lowered_text, cap_map
    
    @staticmethod
    def _restore_capitalization(text: str, cap_map: list) -> str:
        """Restore capitalization to normalized text using position-based mapping.
        
        For each word at position i, if cap_map[i] is True, the first
        alphabetic character of that word is uppercased. This is the
        simple position-based approach used when word counts match.
        
        For across-lines mode where word counts may change significantly,
        use _restore_capitalization_aligned() instead.
        
        Args:
            text: Normalized text (lowercase)
            cap_map: List of booleans from _extract_capitalization
        
        Returns:
            Text with capitalization restored
        """
        words = text.split()
        result = []
        for i, word in enumerate(words):
            if i < len(cap_map) and cap_map[i]:
                result.append(_capitalize_first_alpha(word))
            else:
                result.append(word)
        # If model produced fewer words than original, trailing cap_map
        # entries are simply ignored (no capitalization to apply).
        return ' '.join(result)
    
    @staticmethod
    def _restore_capitalization_aligned(original_text: str, normalized_text: str,
                                         cap_map: list) -> str:
        """Restore capitalization using character-level alignment.
        
        When the model changes word counts (e.g., "inalten" → "in alten"),
        position-based _restore_capitalization() fails because word indices
        shift. This method uses character-level sequence alignment to correctly
        map each normalized word to its corresponding original word, then
        applies capitalization based on the original word's case.
        
        For word splits (multiple normalized words mapping to one original
        word), only the first normalized word is capitalized, matching the
        original word's pattern.
        
        Args:
            original_text: Original text (lowercase, punctuation-stripped;
                           same text that was sent to the model)
            normalized_text: Normalized text (lowercase, without punctuation)
            cap_map: List of booleans from _extract_capitalization,
                     one per word in original_text
        
        Returns:
            Text with capitalization restored
        """
        normalized_words = normalized_text.split()
        if not normalized_words:
            return normalized_text
        
        original_words = original_text.split()
        if not original_words:
            return normalized_text
        
        # If word counts match, use simple position-based restoration
        if len(normalized_words) == len(original_words):
            return MHGNormalizer._restore_capitalization(normalized_text, cap_map)
        
        # Use character-level alignment to map normalized words to original words
        norm_to_orig_word = MHGNormalizer._align_normalized_to_original_charlevel(
            original_text, normalized_text
        )
        
        # Apply capitalization, handling word splits:
        # When multiple normalized words map to the same original word,
        # only the first gets capitalized (matching the original word's
        # pattern where only the first letter is uppercase).
        result = []
        for i, word in enumerate(normalized_words):
            orig_idx = norm_to_orig_word[i]
            if orig_idx < len(cap_map) and cap_map[orig_idx]:
                is_first = (i == 0 or norm_to_orig_word[i - 1] != orig_idx)
                if is_first:
                    result.append(_capitalize_first_alpha(word))
                else:
                    result.append(word)
            else:
                result.append(word)
        
        return ' '.join(result)
    
    @staticmethod
    def _normalize_word_for_alignment(word: str) -> str:
        """Normalize a word for alignment matching.
        
        Converts to lowercase and simplifies common MHG spelling variations
        to improve SequenceMatcher's ability to find correct alignments.
        
        Args:
            word: Word to normalize
        
        Returns:
            Normalized word form for matching
        """
        # Lowercase and long-s to short-s
        result = word.lower().replace('ſ', 's')
        # Remove diacritics for matching
        for char, replacement in [
            ('â', 'a'), ('î', 'i'), ('û', 'u'), ('ô', 'o'),
            ('ä', 'a'), ('ë', 'e'), ('ï', 'i'), ('ö', 'o'), ('ü', 'u'),
            ('é', 'e'), ('è', 'e'), ('ê', 'e'), ('à', 'a'),
            ('â', 'a'), ('î', 'i'), ('ô', 'o'), ('û', 'u'),
        ]:
            result = result.replace(char, replacement)
        return result
    
    @staticmethod
    def _align_normalized_to_original_charlevel(
        original_text: str,
        normalized_text: str
    ) -> List[int]:
        """Align normalized words to original words using word-level sequence matching.
        
        Uses difflib.SequenceMatcher at the word level with normalized word forms
        to find the best alignment. This handles:
        - Word splits (one original → multiple normalized)
        - Word changes (different spelling)
        - Repeated words (matched by sequence context, not just identity)
        
        Args:
            original_text: Original text (space-joined words, punctuation-stripped)
            normalized_text: Normalized text (space-joined words)
        
        Returns:
            List of original word indices, one per normalized word.
            Each normalized word is mapped to the index of the original word
            it aligns to. The mapping is monotonically non-decreasing.
        """
        normalized_words = normalized_text.split()
        if not normalized_words:
            return []
        
        original_words = original_text.split()
        if not original_words:
            return [0] * len(normalized_words)
        
        total_orig = len(original_words)
        total_norm = len(normalized_words)
        
        # If word counts match, use simple position-based mapping
        if total_norm == total_orig:
            return list(range(total_norm))
        
        # Use word-level SequenceMatcher with normalized word forms for better matching
        # This helps match "dise" to "diſe" and "vâhten" to "vochten" as equal blocks
        orig_normalized = [MHGNormalizer._normalize_word_for_alignment(w) for w in original_words]
        norm_normalized = [MHGNormalizer._normalize_word_for_alignment(w) for w in normalized_words]
        
        matcher = difflib.SequenceMatcher(None, orig_normalized, norm_normalized,
                                          autojunk=False)
        
        # Build the mapping from normalized word index to original word index
        norm_to_orig_word = [0] * total_norm
        
        for tag, i1, i2, j1, j2 in matcher.get_opcodes():
            if tag == 'equal':
                # Direct match: normalized words j1:j2 match original words i1:i2
                for k in range(j2 - j1):
                    norm_to_orig_word[j1 + k] = i1 + k
            elif tag == 'replace':
                # Words changed but position preserved (e.g., "vō" → "von").
                # Map proportionally within the block using rounded division so
                # that word splits (e.g., 1 orig → 2 norm) correctly cluster on
                # the same original word rather than drifting to neighbours.
                orig_len = i2 - i1
                norm_len = j2 - j1
                for k in range(norm_len):
                    if orig_len > 0:
                        # Rounded proportional index: maps k ∈ [0, norm_len) to
                        # orig_idx ∈ [i1, i1+orig_len).  When norm_len > orig_len
                        # (word splits), multiple k values round to the same
                        # orig_idx, keeping split fragments together.
                        orig_idx = i1 + min(round(k * orig_len / norm_len), orig_len - 1)
                    else:
                        orig_idx = i1
                    norm_to_orig_word[j1 + k] = min(orig_idx, total_orig - 1)
            elif tag == 'insert':
                # Extra words in normalized (e.g., from word splits).
                # Map to the next original word position (i1), since inserts in
                # normalization are typically split fragments that belong to the
                # following original word, not the preceding one.
                for j in range(j1, j2):
                    if i1 < total_orig:
                        norm_to_orig_word[j] = i1
                    else:
                        norm_to_orig_word[j] = total_orig - 1
            # 'delete': original words removed, no normalized words to map
        
        # Ensure monotonicity: each normalized word should map to an original
        # word at or after the previous normalized word's mapping
        prev_orig = 0
        for i in range(total_norm):
            if norm_to_orig_word[i] < prev_orig:
                norm_to_orig_word[i] = prev_orig
            prev_orig = norm_to_orig_word[i]
        
        return norm_to_orig_word
    
    @staticmethod
    def _spread_repeated_mappings(
        norm_to_orig_word: List[int],
        original_words: List[str],
        normalized_words: List[str]
    ) -> List[int]:
        """Spread repeated word mappings across their occurrences in the original.
        
        When the character-level alignment maps multiple identical normalized words
        to the same original position, this post-processing step spreads them across
        the actual occurrences of that word in the original text.
        
        For example, if "dise" appears 6 times in the original (at indices 0, 4, 8, 12, 16, 20)
        and 6 normalized "dise" words all map to orig[0], this spreads them to
        [0, 4, 8, 12, 16, 20].
        
        Args:
            norm_to_orig_word: Initial alignment mapping
            original_words: List of original words
            normalized_words: List of normalized words
        
        Returns:
            Adjusted alignment mapping with repeated words spread across occurrences
        """
        if not norm_to_orig_word:
            return norm_to_orig_word
        
        # Build a map of normalized word -> list of original indices where it appears
        orig_word_occurrences = {}
        for i, word in enumerate(original_words):
            word_lower = word.lower()
            if word_lower not in orig_word_occurrences:
                orig_word_occurrences[word_lower] = []
            orig_word_occurrences[word_lower].append(i)
        
        # Track usage of each original word occurrence
        orig_usage = {}  # orig_idx -> count of times used
        
        # Track which normalized word forms we've seen and how many times
        norm_word_count = {}  # normalized_word -> count seen so far
        
        result = []
        for i, orig_idx in enumerate(norm_to_orig_word):
            norm_word_lower = normalized_words[i].lower()
            
            # Check if this normalized word form has multiple occurrences in original
            occurrences = orig_word_occurrences.get(norm_word_lower, [])
            
            if len(occurrences) > 1:
                # This word appears multiple times in original
                # Use the next available occurrence
                count_so_far = norm_word_count.get(norm_word_lower, 0)
                norm_word_count[norm_word_lower] = count_so_far + 1
                
                if count_so_far < len(occurrences):
                    # Map to the next occurrence
                    new_orig_idx = occurrences[count_so_far]
                    result.append(new_orig_idx)
                    orig_usage[new_orig_idx] = orig_usage.get(new_orig_idx, 0) + 1
                else:
                    # More occurrences in normalized than original; use alignment
                    result.append(orig_idx)
                    orig_usage[orig_idx] = orig_usage.get(orig_idx, 0) + 1
            else:
                # Unique word or single occurrence; use alignment
                result.append(orig_idx)
                orig_usage[orig_idx] = orig_usage.get(orig_idx, 0) + 1
        
        return result

    @staticmethod
    def _assign_lines_by_word_alignment(
        original_lines: List[str],
        normalized_text: str,
    ) -> List[int]:
        """Assign each word in normalized_text to a line index using word-level alignment.

        Uses _align_normalized_to_original_charlevel (word-level difflib with
        normalized word forms) to map each normalized word to its corresponding
        original word index, then looks up which line that original word
        belongs to.

        This approach is robust to MHG character differences (ſ→s, diacritics,
        etc.) because the alignment normalizes word forms before matching.
        It also handles word splits correctly — multiple normalized words
        mapping to the same original word stay on the same line.
        """
        normalized_words = normalized_text.split()
        if not normalized_words:
            return []

        num_lines = len(original_lines)
        full_orig = ' '.join(original_lines)

        # Get word-level alignment: each normalized word → original word index
        norm_to_orig_word = MHGNormalizer._align_normalized_to_original_charlevel(
            full_orig, normalized_text
        )

        # Build original word index → line mapping
        word_to_line = []
        for line_idx, line_text in enumerate(original_lines):
            for _ in line_text.split():
                word_to_line.append(line_idx)

        # Map each normalized word to its line via the original word index
        norm_word_to_line = []
        prev_line = 0
        for orig_word_idx in norm_to_orig_word:
            if orig_word_idx < len(word_to_line):
                line_idx = word_to_line[orig_word_idx]
            else:
                line_idx = prev_line
            norm_word_to_line.append(line_idx)
            prev_line = line_idx

        # Post-processing: fix cross-line split fragments.
        # When a word at a line boundary (e.g., "en" from "Engülden" → "en gültin")
        # was incorrectly assigned to the previous line, detect it by checking if
        # the normalized word is a prefix of the first word of the next original line.
        total_norm = len(normalized_words)
        if len(original_lines) > 1:
            for norm_idx in range(total_norm - 1):
                curr_line = norm_word_to_line[norm_idx]
                next_line = norm_word_to_line[norm_idx + 1]

                # Only care about descending line assignments (word drifted to previous line)
                if curr_line < next_line and curr_line < num_lines - 1:
                    # This word is on an earlier line than the next word.
                    # Check if it's a split fragment: does the next original line's
                    # first word start with this normalized word (after normalizing)?
                    next_orig_line = curr_line + 1
                    next_orig_words = original_lines[next_orig_line].split()
                    if next_orig_words:
                        first_next_word = next_orig_words[0]
                        norm_current = MHGNormalizer._normalize_word_for_alignment(normalized_words[norm_idx])
                        norm_first_next = MHGNormalizer._normalize_word_for_alignment(first_next_word)

                        # If the current word is a prefix of the first word of the next original line,
                        # it's likely a split fragment that belongs on the next line
                        if (norm_current and
                                norm_first_next.startswith(norm_current) and
                                len(norm_current) <= 4):  # short fragments are split candidates
                            norm_word_to_line[norm_idx] = next_line

        return norm_word_to_line

    @staticmethod
    def _get_normalized_to_original_char_mapping(
        original_text: str,
        normalized_text: str
    ) -> dict:
        """Map each normalized word start position to its corresponding original character position.
        
        Uses character-level SequenceMatcher to align normalized text to original text,
        then computes the starting character position of each normalized word and maps
        it to the corresponding position in the original text.
        
        This mapping is used to determine which line each normalized word belongs to
        when splitting the normalized output back into lines. Unlike word-level alignment,
        character-level alignment correctly handles cases where the model merges multiple
        short lines into one.
        
        Args:
            original_text: Original text (space-joined words, may contain newlines)
            normalized_text: Normalized text (space-joined words)
        
        Returns:
            Dictionary mapping normalized word index to original character position
        """
        normalized_words = normalized_text.split()
        if not normalized_words:
            return {}
        
        # Use character-level SequenceMatcher to align normalized text to original text
        matcher = difflib.SequenceMatcher(None, original_text, normalized_text,
                                          autojunk=False)
        
        # Build character position mapping from normalized to original
        norm_to_orig_char_pos = {}
        for tag, i1, i2, j1, j2 in matcher.get_opcodes():
            if tag == 'equal':
                for k in range(j2 - j1):
                    norm_to_orig_char_pos[j1 + k] = i1 + k
            elif tag == 'replace':
                orig_len = i2 - i1
                norm_len = j2 - j1
                if orig_len > 0 and norm_len > 0:
                    for k in range(norm_len):
                        norm_to_orig_char_pos[j1 + k] = i1 + min(
                            k * orig_len // norm_len, orig_len - 1)
                elif norm_len > 0:
                    for k in range(norm_len):
                        norm_to_orig_char_pos[j1 + k] = i1 if i1 < len(original_text) else len(original_text) - 1
            elif tag == 'insert':
                for j in range(j1, j2):
                    if i1 > 0:
                        norm_to_orig_char_pos[j] = i1 - 1
                    elif i1 < len(original_text):
                        norm_to_orig_char_pos[j] = i1
        
        # Compute start position of each normalized word
        word_start_positions = []
        pos = 0
        for word in normalized_words:
            word_start_positions.append(pos)
            pos += len(word) + 1  # +1 for space
        
        # Map each normalized word to its original character position
        word_to_orig_char = {}
        for i, start_pos in enumerate(word_start_positions):
            if start_pos in norm_to_orig_char_pos:
                word_to_orig_char[i] = norm_to_orig_char_pos[start_pos]
            else:
                # Fallback: find nearest mapped position
                for p in range(start_pos, -1, -1):
                    if p in norm_to_orig_char_pos:
                        word_to_orig_char[i] = norm_to_orig_char_pos[p]
                        break
                else:
                    word_to_orig_char[i] = 0
        
        return word_to_orig_char
    
    @staticmethod
    def _align_normalized_to_original(
        original_words: List[str],
        normalized_words: List[str]
    ) -> List[int]:
        """Align normalized words to original words using word-level sequence matching.
        
        .. deprecated::
            Use :meth:`_align_normalized_to_original_charlevel` instead, which
            uses character-level alignment and correctly handles word splits at
            line boundaries. This word-level method can incorrectly match a
            split word fragment (e.g., "in" from "inallen" → "in allen") to a
            standalone identical word on a different line.
        
        Uses difflib.SequenceMatcher to map each normalized word to its
        corresponding original word index. This handles word splits and merges
        within a single line reasonably well, but can produce incorrect line
        assignments when a split fragment matches a word on a different line.
        
        Args:
            original_words: Words from the joined original text
            normalized_words: Words from the model's normalized output
        
        Returns:
            List of original word indices, one per normalized word.
            Each normalized word is mapped to the index of the original word
            it aligns to. For insertions (extra words from splits), the
            preceding original word's index is used.
        """
        if not normalized_words:
            return []
        
        if not original_words:
            return [0] * len(normalized_words)
        
        # Use SequenceMatcher to align normalized words to original words.
        # autojunk=False ensures common words (like "und", "der") are still
        # used for anchoring the alignment, which improves quality.
        matcher = difflib.SequenceMatcher(None, original_words, normalized_words,
                                          autojunk=False)
        
        # Map each normalized word index to its corresponding original word index
        norm_to_orig = [0] * len(normalized_words)
        
        for tag, i1, i2, j1, j2 in matcher.get_opcodes():
            if tag == 'equal':
                # Direct match: normalized word j corresponds to original word i
                for k in range(j2 - j1):
                    norm_to_orig[j1 + k] = i1 + k
            elif tag == 'replace':
                # Words changed but position preserved (e.g., "vō" → "von").
                # Map each normalized word to the corresponding original word
                # by proportional position within the replace block.
                orig_len = i2 - i1
                norm_len = j2 - j1
                for k in range(norm_len):
                    orig_idx = i1 + min(k * orig_len // norm_len, orig_len - 1) if norm_len > 0 else i1
                    norm_to_orig[j1 + k] = orig_idx
            elif tag == 'insert':
                # Extra words in normalized text (e.g., from word splits like
                # "inalten" → "in" "alten"). Assign to the original word just
                # before the insertion point.
                for j in range(j1, j2):
                    if i1 > 0:
                        norm_to_orig[j] = i1 - 1
                    elif i1 < len(original_words):
                        norm_to_orig[j] = i1
            # 'delete': words removed from original, no normalized words to assign
        
        return norm_to_orig
    
    def _get_token_count(self, text: str) -> int:
        """
        Get the number of tokens for a text including the 'normalize: ' prefix.
        
        Args:
            text: Input text (without prefix)
        
        Returns:
            Total token count including prefix and special tokens
        """
        prefix = "normalize: "
        input_text = prefix + text
        return len(self.tokenizer.encode(input_text))
    
    def _create_word_windows(
        self,
        words: List[str],
        max_window_tokens: int,
        stride_step_tokens: int
    ) -> List[Tuple[int, int]]:
        """
        Create overlapping word windows that fit within the token budget.
        
        Each window is defined by (start_word_idx, end_word_idx) where
        end_word_idx is exclusive. Windows overlap so that every word
        appears in the interior of at least one window, ensuring good
        context for the normalization of each word.
        
        Args:
            words: List of words from the input text
            max_window_tokens: Maximum number of tokens per window
                               (including the 'normalize: ' prefix)
            stride_step_tokens: Number of tokens to advance between
                                consecutive windows
        
        Returns:
            List of (start_idx, end_idx) tuples defining word windows
        """
        if not words:
            return []
        
        prefix = "normalize: "
        prefix_tokens = len(self.tokenizer.encode(prefix, add_special_tokens=False))
        effective_max = max_window_tokens - prefix_tokens
        
        if effective_max <= 0:
            # Window too small for any content, fall back to single-word windows
            return [(i, i + 1) for i in range(len(words))]
        
        # Pre-compute token count for each word (with leading space for joining)
        word_token_counts = []
        for w in words:
            tc = len(self.tokenizer.encode(" " + w, add_special_tokens=False))
            word_token_counts.append(tc)
        
        # Cumulative token counts for efficient range queries
        cum_tokens = [0]
        for tc in word_token_counts:
            cum_tokens.append(cum_tokens[-1] + tc)
        
        windows = []
        start_idx = 0
        
        while start_idx < len(words):
            # Find the end of this window: the largest end_idx where
            # the token count from start_idx to end_idx fits in effective_max
            end_idx = start_idx
            while end_idx < len(words) and cum_tokens[end_idx + 1] - cum_tokens[start_idx] <= effective_max:
                end_idx += 1
            
            # Ensure at least one word per window (even if it exceeds token limit)
            if end_idx == start_idx:
                end_idx = start_idx + 1
            
            windows.append((start_idx, end_idx))
            
            # If we've covered all words, stop
            if end_idx >= len(words):
                break
            
            # Advance start by stride_step_tokens worth of words
            new_start = start_idx
            while new_start < end_idx and cum_tokens[new_start + 1] - cum_tokens[start_idx] < stride_step_tokens:
                new_start += 1
            
            # Ensure we make progress (at least one word advance)
            if new_start <= start_idx:
                new_start = start_idx + 1
            
            start_idx = new_start
        
        return windows
    
    def _merge_window_predictions(
        self,
        total_words: int,
        windows: List[Tuple[int, int]],
        window_predictions: List[List[str]],
        original_words: List[str]
    ) -> List[str]:
        """
        Merge predictions from overlapping windows, preferring window centers.
        
        For each word position, the prediction from the window where that
        word is closest to the center is used. This ensures that words in
        overlap regions get the best context-informed prediction, since
        the center of a window has the most surrounding context.
        
        When the model produces a different number of words than expected
        (due to word splits or merges), sequence alignment is used instead
        of proportional mapping to prevent words from disappearing.
        
        Args:
            total_words: Total number of words in the original text
            windows: List of (start_idx, end_idx) tuples defining word windows
            window_predictions: List of word lists, one per window
            original_words: Original input words (used as fallback)
        
        Returns:
            List of merged prediction words
        """
        # Each position stores a list of words (to handle splits where
        # multiple predicted words map to one original position).
        merged = [None] * total_words
        best_distances = [float('inf')] * total_words
        
        for (start_idx, end_idx), pred_words in zip(windows, window_predictions):
            # Center of the window
            center = (start_idx + end_idx - 1) / 2.0
            expected_len = end_idx - start_idx
            
            if len(pred_words) == expected_len:
                # Perfect alignment: direct word-by-word mapping
                for i in range(expected_len):
                    word_pos = start_idx + i
                    if word_pos < total_words:
                        distance = abs(word_pos - center)
                        if distance < best_distances[word_pos]:
                            best_distances[word_pos] = distance
                            merged[word_pos] = [pred_words[i]]
            else:
                # Length mismatch: the model produced a different number of words
                # than expected (word splits or merges). Use character-level
                # alignment to properly map each predicted word to its original
                # position, preventing word disappearance and correctly handling
                # word splits (e.g., "inalten" → "in alten").
                window_words = original_words[start_idx:end_idx]
                window_text = ' '.join(window_words)
                pred_text = ' '.join(pred_words)
                alignment = self._align_normalized_to_original_charlevel(
                    window_text, pred_text
                )
                
                # Group predicted words by their aligned original position.
                # Multiple predicted words mapping to the same position
                # (e.g., from a split like "inalten" → "in" "alten") are
                # collected together so no words are lost.
                pos_to_preds = {}
                for j, orig_idx in enumerate(alignment):
                    abs_pos = start_idx + orig_idx
                    if abs_pos not in pos_to_preds:
                        pos_to_preds[abs_pos] = []
                    pos_to_preds[abs_pos].append(pred_words[j])
                
                for abs_pos, preds in pos_to_preds.items():
                    if abs_pos < total_words:
                        distance = abs(abs_pos - center)
                        if distance < best_distances[abs_pos]:
                            best_distances[abs_pos] = distance
                            merged[abs_pos] = preds
        
        # Fill any unfilled positions with original words (model deletions
        # within a window — preserve the original un-normalized word rather
        # than dropping it entirely).
        for i in range(total_words):
            if merged[i] is None:
                merged[i] = [original_words[i]]
        
        # Flatten lists of words into a single word list.
        # Splits produce multiple words at one position; merges produce one.
        result = []
        for words_at_pos in merged:
            result.extend(words_at_pos)
        
        return result
    
    def normalize_long_text(
        self,
        text: str,
        max_length: Optional[int] = None,
        stride_window_tokens: Optional[int] = None,
        stride_step_tokens: Optional[int] = None,
        num_beams: Optional[int] = None,
        length_penalty: Optional[float] = None,
        do_sample: Optional[bool] = None,
        top_p: Optional[float] = None,
        temperature: Optional[float] = None,
        preserve_punctuation: Optional[bool] = None,
        preserve_capitalization: Optional[bool] = None,
        progress_callback: Optional[callable] = None
    ) -> str:
        """
        Normalize a potentially long text using sliding window with stride.
        
        Splits the input text into overlapping windows of words, normalizes
        each window independently, then merges the predictions. For each word
        position, the prediction from the window where that word is closest
        to the center is preferred, as center predictions have the most
        surrounding context.
        
        If the text fits within max_length tokens, delegates to normalize_text()
        for single-pass normalization (no windowing needed).
        
        Args:
            text: Input text to normalize (original MHG)
            max_length: Maximum generation length (tokens)
            stride_window_tokens: Max input tokens per window (including prefix).
                                  Default: 400 (from config)
            stride_step_tokens: Token advance between windows. Controls overlap;
                                overlap = window - step. Default: 300 (from config)
            num_beams: Number of beams for beam search
            length_penalty: Length penalty for beam search
            do_sample: Whether to use sampling
            top_p: Nucleus sampling threshold
            temperature: Temperature for sampling
            preserve_punctuation: If True, strip editorial punctuation before
                                  normalization and restore it afterwards.
                                  Default: None (read from config, falls back to True)
            preserve_capitalization: If True, record capitalization before
                                     normalization (lowercasing the input) and
                                     restore it afterwards.
                                     Default: None (read from config, falls back to True)
        
        Returns:
            Normalized text
        """
        # Resolve parameters
        max_length = max_length or self.config['inference']['max_length']
        stride_window_tokens = stride_window_tokens or self.config['inference'].get('stride_window_tokens', 400)
        stride_step_tokens = stride_step_tokens or self.config['inference'].get('stride_step_tokens', 300)
        if preserve_punctuation is None:
            preserve_punctuation = self.config['inference'].get('preserve_punctuation', True)
        if preserve_capitalization is None:
            preserve_capitalization = self.config['inference'].get('preserve_capitalization', True)
        
        # Extract editorial punctuation before normalization
        punct_map = None
        if preserve_punctuation:
            text, punct_map = self._extract_punctuation(text)
        
        # Extract capitalization before normalization (after punctuation extraction
        # so cap_map indices align with punct_map indices)
        cap_map = None
        if preserve_capitalization:
            text, cap_map = self._extract_capitalization(text)
        
        # Check if stride is actually needed
        token_count = self._get_token_count(text)
        if token_count <= max_length:
            # Text fits in a single window, use simple normalization
            result = self.normalize_text(
                text, max_length=max_length, stride=False,
                num_beams=num_beams, length_penalty=length_penalty,
                do_sample=do_sample, top_p=top_p, temperature=temperature,
                preserve_punctuation=False,  # already stripped
                preserve_capitalization=False  # already lowercased
            )
            if cap_map is not None:
                result = self._restore_capitalization(result, cap_map)
            if punct_map is not None:
                result = self._restore_punctuation(result, punct_map)
            return result
        
        # Split text into words
        words = text.split()
        if not words:
            return ""
        
        if len(words) == 1:
            # Single word, just normalize directly (even if token count exceeds limit)
            result = self.normalize_text(
                text, max_length=max_length, stride=False,
                num_beams=num_beams, length_penalty=length_penalty,
                do_sample=do_sample, top_p=top_p, temperature=temperature,
                preserve_punctuation=False,  # already stripped
                preserve_capitalization=False  # already lowercased
            )
            if cap_map is not None:
                result = self._restore_capitalization(result, cap_map)
            if punct_map is not None:
                result = self._restore_punctuation(result, punct_map)
            return result
        
        # Create overlapping windows
        windows = self._create_word_windows(words, stride_window_tokens, stride_step_tokens)
        
        num_windows = len(windows)
        # if num_windows > 1:
        #     print(f"  Stride: {token_count} tokens > {max_length} max, "
        #           f"using {num_windows} windows "
        #           f"(window={stride_window_tokens}tok, step={stride_step_tokens}tok, "
        #           f"overlap={stride_window_tokens - stride_step_tokens}tok)")
        
        # Normalize each window
        window_predictions = []
        windows_iter = windows
        if num_windows > 1:
            windows_iter = tqdm(windows, desc="  Progress",
                               total=num_windows, ncols=100,
                               bar_format='{desc}: {percentage:3.0f}%|{bar}| {elapsed}')
        for wi, (start_idx, end_idx) in enumerate(windows_iter):
            window_text = ' '.join(words[start_idx:end_idx])
            prediction = self.normalize_text(
                window_text, max_length=max_length, stride=False,
                num_beams=num_beams, length_penalty=length_penalty,
                do_sample=do_sample, top_p=top_p, temperature=temperature,
                preserve_punctuation=False,  # already stripped
                preserve_capitalization=False  # already lowercased
            )
            # Report progress scaled to a low range (5-75%) so the UI bar
            # reflects the fact that significant post-processing work remains
            # after this method returns.  With 5 beams, each window is
            # expensive, but the caller-side work (character-level alignment
            # in _align_normalized_to_original_charlevel, line reassignment,
            # cap/punct restoration in normalize_multiline_text) can easily
            # take as long as or longer than the window loop itself.
            #
            # Scale: window 0 → 5%, last window → 75%.
            if progress_callback:
                pct = 5 + int(wi * 70 / max(num_windows - 1, 1))
                progress_callback(pct, 100)
            pred_words = prediction.split()
            window_predictions.append(pred_words)
        
        # Merge predictions from overlapping windows
        merged_words = self._merge_window_predictions(
            len(words), windows, window_predictions, words
        )
        
        # Signal that generation + merge are done.  The remaining 25% of
        # progress will be consumed by caller-side post-processing
        # (alignment, capitalization/punctuation restoration, line splitting)
        # before app.py sets progress to 100.
        if progress_callback:
            progress_callback(78, 100)
        
        result = ' '.join(merged_words)
        if cap_map is not None:
            # Use character-level alignment for capitalization restoration
            # because word counts may change due to splits/merges.
            result = self._restore_capitalization_aligned(
                text, result, cap_map
            )
        if punct_map is not None:
            # Use character-level alignment for punctuation restoration
            # instead of position-based mapping, because word counts may
            # change due to splits/merges in the merge step.
            result = self._restore_punctuation_aligned(
                text, result, punct_map
            )
        return result
    
    def normalize_text(
        self,
        text: str,
        max_length: Optional[int] = None,
        num_beams: Optional[int] = None,
        length_penalty: Optional[float] = None,
        do_sample: Optional[bool] = None,
        top_p: Optional[float] = None,
        temperature: Optional[float] = None,
        stride: Optional[bool] = None,
        stride_window_tokens: Optional[int] = None,
        stride_step_tokens: Optional[int] = None,
        preserve_punctuation: Optional[bool] = None,
        preserve_capitalization: Optional[bool] = None,
        attach_en_proclitic: Optional[bool] = None,
        lenition_t_after_n: Optional[bool] = None,
        lenition_t_after_l: Optional[bool] = None,
        niet_to_niht: Optional[bool] = None,
        common_apocopes: Optional[bool] = None,
        progress_callback: Optional[callable] = None
    ) -> str:
        """
        Normalize a single text string.
        
        When stride=True (default) and the input text exceeds the model's
        maximum sequence length, automatically uses sliding-window inference
        via normalize_long_text() to process the full text without truncation.
        
        Args:
            text: Input text to normalize (original MHG)
            max_length: Maximum generation length
            num_beams: Number of beams for beam search
            length_penalty: Length penalty for beam search
            do_sample: Whether to use sampling (vs beam search)
            top_p: Nucleus sampling threshold (0.0 to 1.0)
            temperature: Temperature for sampling (0.1 to 2.0)
            stride: Whether to use sliding window for long texts.
                    Default: True (from config). Set to False to always truncate.
            stride_window_tokens: Max input tokens per window (for stride mode)
            stride_step_tokens: Token advance between windows (for stride mode)
            preserve_punctuation: If True, strip editorial punctuation before
                                  normalization and restore it afterwards.
                                  Default: None (read from config, falls back to True)
            preserve_capitalization: If True, record capitalization before
                                     normalization (lowercasing the input) and
                                     restore it afterwards.
                                     Default: None (read from config, falls back to True)
            attach_en_proclitic: If True, attach free-standing negation proclitic
                                  "en" to the following word with a hyphen
                                  (e.g. "en guot" → "en-guot"). Applied as the
                                  final post-processing step.
                                  Default: None (read from config, falls back to False)
        
        Returns:
            Normalized text
        """
        # Resolve preserve_punctuation
        if preserve_punctuation is None:
            preserve_punctuation = self.config['inference'].get('preserve_punctuation', True)
        
        # Resolve preserve_capitalization
        if preserve_capitalization is None:
            preserve_capitalization = self.config['inference'].get('preserve_capitalization', True)

        # Resolve post-processing options
        if attach_en_proclitic is None:
            attach_en_proclitic = self.config.get('post_processing', {}).get('attach_en_proclitic', False)
        if lenition_t_after_n is None:
            lenition_t_after_n = self.config.get('post_processing', {}).get('lenition_t_after_n', False)
        if lenition_t_after_l is None:
            lenition_t_after_l = self.config.get('post_processing', {}).get('lenition_t_after_l', False)
        if niet_to_niht is None:
            niet_to_niht = self.config.get('post_processing', {}).get('niet_to_niht', False)
        if common_apocopes is None:
            common_apocopes = self.config.get('post_processing', {}).get('common_apocopes', False)
        
        # Extract editorial punctuation before normalization
        punct_map = None
        if preserve_punctuation:
            text, punct_map = self._extract_punctuation(text)
        
        # Extract capitalization before normalization (after punctuation extraction
        # so cap_map indices align with punct_map indices)
        cap_map = None
        if preserve_capitalization:
            text, cap_map = self._extract_capitalization(text)
        
        # Resolve stride setting
        if stride is None:
            stride = self.config['inference'].get('stride', True)
        
        # If stride is enabled, check if we need it
        if stride:
            max_length_resolved = max_length or self.config['inference']['max_length']
            token_count = self._get_token_count(text)
            if token_count > max_length_resolved:
                result = self.normalize_long_text(
                    text, max_length=max_length_resolved,
                    stride_window_tokens=stride_window_tokens,
                    stride_step_tokens=stride_step_tokens,
                    num_beams=num_beams, length_penalty=length_penalty,
                    do_sample=do_sample, top_p=top_p, temperature=temperature,
                    preserve_punctuation=False,  # already stripped above
                    preserve_capitalization=False,  # already lowercased above
                    progress_callback=progress_callback
                )
                if cap_map is not None:
                    result = self._restore_capitalization(result, cap_map)
                if punct_map is not None:
                    result = self._restore_punctuation(result, punct_map)
                # Post-processing: common apocopes (before en-attachment so "en wile" → "en wil" → "enwil")
                if common_apocopes:
                    result = _common_apocopes(result)
                # Post-processing: attach negation proclitic
                if attach_en_proclitic:
                    result = _attach_en_proclitic(result)
                # Post-processing: lenition of t after n before vowel
                if lenition_t_after_n:
                    result = _lenition_t_after_n_before_vowel(result)
                # Post-processing: lenition of t after l before vowel
                if lenition_t_after_l:
                    result = _lenition_t_after_l_before_vowel(result)
                # Post-processing: niet → niht
                if niet_to_niht:
                    result = _niet_to_niht(result)
                return result
        
        # Use config values if not provided
        max_length = max_length or self.config['inference']['max_length']
        num_beams = num_beams or self.config['inference']['num_beams']
        length_penalty = length_penalty or self.config['inference']['length_penalty']
        do_sample = do_sample if do_sample is not None else self.config['inference']['do_sample']
        top_p = top_p if top_p is not None else self.config['inference']['top_p']
        temperature = temperature if temperature is not None else self.config['inference'].get('temperature', 1.0)
        
        # Add prefix if using T5-style model
        prefix = "normalize: "
        input_text = prefix + text
        
        # Tokenize
        inputs = self.tokenizer(
            input_text,
            return_tensors='pt',
            padding=True,
            truncation=True,
            max_length=max_length
        )
        
        # Move to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Generate
        with torch.no_grad():
            # Prepare generation arguments by unpacking inputs directly
            generate_kwargs = {
                **inputs,  # This includes 'input_ids' and 'attention_mask'
                'max_length': max_length,
                'num_beams': num_beams,
                'length_penalty': length_penalty,
                'do_sample': do_sample,
                'early_stopping': True
            }
            # Only add sampling-specific parameters if sampling is enabled
            if do_sample:
                generate_kwargs['top_p'] = top_p
                generate_kwargs['temperature'] = temperature
            
            outputs = self.model.generate(**generate_kwargs)
        
        # Decode
        normalized = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Restore capitalization (before punctuation, since cap_map indices
        # align with punctuation-stripped words)
        if cap_map is not None:
            normalized = self._restore_capitalization(normalized, cap_map)
        
        # Restore editorial punctuation
        if punct_map is not None:
            normalized = self._restore_punctuation(normalized, punct_map)

        # Post-processing: common apocopes (before en-attachment so "en wile" → "en wil" → "enwil")
        if common_apocopes:
            normalized = _common_apocopes(normalized)
        # Post-processing: attach negation proclitic (after apocopes)
        if attach_en_proclitic:
            normalized = _attach_en_proclitic(normalized)
        # Post-processing: lenition of t after n before vowel
        if lenition_t_after_n:
            normalized = _lenition_t_after_n_before_vowel(normalized)
        # Post-processing: lenition of t after l before vowel
        if lenition_t_after_l:
            normalized = _lenition_t_after_l_before_vowel(normalized)
        # Post-processing: niet → niht
        if niet_to_niht:
            normalized = _niet_to_niht(normalized)

        return normalized
    
    def normalize_texts(
        self,
        texts: List[str],
        batch_size: int = 8,
        show_progress: bool = True,
        **kwargs
    ) -> List[str]:
        """
        Normalize multiple texts efficiently using batching.
        
        Args:
            texts: List of input texts
            batch_size: Batch size for processing
            show_progress: Whether to show progress bar
            **kwargs: Additional arguments for normalize_text
                      (including stride, stride_window_tokens, stride_step_tokens)
        
        Returns:
            List of normalized texts
        """
        normalized_texts = []
        
        # Calculate total number of batches
        total_batches = (len(texts) + batch_size - 1) // batch_size
        
        # Create progress bar if requested
        iterator = range(0, len(texts), batch_size)
        if show_progress:
            iterator = tqdm(iterator, desc="Normalizing lines",
                          total=total_batches,
                          ncols=100,
                          bar_format='{desc}: {percentage:3.0f}%|{bar}| {elapsed}')
        
        for i in iterator:
            batch = texts[i:i + batch_size]
            batch_normalized = [
                self.normalize_text(text, **kwargs)
                for text in batch
            ]
            normalized_texts.extend(batch_normalized)
        
        return normalized_texts
    
    def normalize_file(
        self,
        input_path: str,
        output_path: str,
        batch_size: int = 8,
        across_lines: bool = True,
        preserve_punctuation: Optional[bool] = None,
        preserve_capitalization: Optional[bool] = None,
        **kwargs
    ):
        """
        Normalize texts from a file.
        
        When across_lines=True (default), all lines are joined into a single
        continuous text and normalized together, utilizing the model's full
        capacity and sliding-window stride across line boundaries. The output
        is split back into lines based on original word counts per line.
        
        When across_lines=False, each line is normalized independently
        (original behavior).
        
        Args:
            input_path: Path to input file (one text per line)
            output_path: Path to output file
            batch_size: Batch size for processing (only used when across_lines=False)
            across_lines: If True, join all lines and normalize as one continuous
                          text, then split back by word counts. Default: True.
            preserve_punctuation: If True, strip editorial punctuation before
                                  normalization and restore it afterwards.
                                  Default: None (read from config, falls back to True)
            preserve_capitalization: If True, record capitalization before
                                     normalization (lowercasing the input) and
                                     restore it afterwards.
                                     Default: None (read from config, falls back to True)
            **kwargs: Additional arguments for normalize_text
                      (including stride, stride_window_tokens, stride_step_tokens)
        """
        input_path = Path(input_path)
        output_path = Path(output_path)
        
        # Resolve preserve_punctuation
        if preserve_punctuation is None:
            preserve_punctuation = self.config['inference'].get('preserve_punctuation', True)
        
        # Resolve preserve_capitalization
        if preserve_capitalization is None:
            preserve_capitalization = self.config['inference'].get('preserve_capitalization', True)
        
        # Read input file — do NOT lowercase here; capitalization extraction
        # handles lowercasing when preserve_capitalization is True, otherwise
        # we lowercase explicitly below.
        # Preserve empty lines: track their positions so they can be reinserted
        # in the output at the same locations.
        print(f"\nNormalizing: {input_path}")
        with open(input_path, 'r', encoding='utf-8') as f:
            all_lines = [line.rstrip('\n').rstrip('\r') for line in f]
        
        # Build a map of which line indices are empty, and extract non-empty
        # lines for normalization. Empty lines will be reinserted as-is.
        empty_line_indices = set()
        texts = []
        for i, line in enumerate(all_lines):
            if line.strip() == '':
                empty_line_indices.add(i)
            else:
                texts.append(line.strip())
        
        # Map from non-empty line index to original line index, so we can
        # place normalized output back into the correct positions.
        non_empty_to_original = []
        for i, line in enumerate(all_lines):
            if line.strip() != '':
                non_empty_to_original.append(i)
        
        # If preserve_capitalization is disabled, lowercase all input now
        # (the model expects lowercase input). When preserve_capitalization
        # is enabled, _extract_capitalization will lowercase the text and
        # record the original capitalization pattern for later restoration.
        if not preserve_capitalization:
            texts = [line.lower() for line in texts]
        
        total_texts = len(texts)
        #print(f"Loaded {total_texts} lines")
        
        if across_lines:
            # Per-line pre-processing: extract punctuation and capitalization
            # before joining into a single text for normalization.
            # Order matters: punctuation extraction first (on original-case text),
            # then capitalization extraction (on punctuation-stripped text).
            # This ensures cap_map and punct_map are both indexed by the same
            # word list (punctuation-stripped words), making per-line restoration
            # straightforward.
            line_punct_maps = None
            line_cap_maps = None
            stripped_lines = None   # punctuation-stripped lines (original case)
            lowered_lines = None    # lowercase, punctuation-stripped lines
            
            if preserve_punctuation:
                line_punct_maps = []
                stripped_lines = []
                for line in texts:
                    stripped_line, punct_map = self._extract_punctuation(line)
                    stripped_lines.append(stripped_line)
                    line_punct_maps.append(punct_map)
            
            if preserve_capitalization:
                line_cap_maps = []
                lowered_lines = []
                # Capitalization extraction operates on the punctuation-stripped
                # text (if punctuation was extracted) or the original text.
                cap_source_lines = stripped_lines if stripped_lines is not None else texts
                for line in cap_source_lines:
                    lowered_line, cap_map = self._extract_capitalization(line)
                    lowered_lines.append(lowered_line)
                    line_cap_maps.append(cap_map)
            
            # Determine the text to send to the model:
            # - If both preserved: lowered_lines (lowercase, no punctuation)
            # - If only cap preserved: lowered_lines (lowercase, with punctuation)
            # - If only punct preserved: stripped_lines (original case, no punctuation)
            # - If neither preserved: texts (lowercased earlier)
            if lowered_lines is not None:
                word_counts = [len(line.split()) for line in lowered_lines]
                full_text = ' '.join(lowered_lines)
            elif stripped_lines is not None:
                word_counts = [len(line.split()) for line in stripped_lines]
                full_text = ' '.join(stripped_lines)
            else:
                word_counts = [len(line.split()) for line in texts]
                full_text = ' '.join(texts)
            
            normalized_full = self.normalize_text(
                full_text, preserve_punctuation=False,
                preserve_capitalization=False, **kwargs
            )
            
            # Split normalized output back into lines using occurrence-based tracking
            # with word split detection.
            #
            # This correctly handles:
            # - Word splits (e.g., "enwere" → "en wære"): detected via alignment
            # - Repeated words (e.g., "dise" x10): occurrence tracking maps k-th occurrence
            #   in normalized to k-th occurrence in original
            # - Different word counts: proportional mapping fallback
            num_lines = len(texts)
            normalized_words = normalized_full.split()
            total_norm = len(normalized_words)
            
            # Build word_to_line mapping from original word indices to line indices
            word_to_line = []
            for line_idx, count in enumerate(word_counts):
                for _ in range(count):
                    word_to_line.append(line_idx)
            
            # Calculate total original words
            total_orig = sum(word_counts)
            
            # Build original word list for occurrence tracking
            original_words = full_text.split()
            
            # Helper function to normalize word for occurrence tracking
            def normalize_word_for_tracking(word: str) -> str:
                """Normalize word for occurrence tracking (case-insensitive, ſ→s, simplified)."""
                # Basic normalization: lowercase and long-s to short-s
                base = word.lower().replace('ſ', 's')
                # Simplify common MHG spelling variations for matching
                # Remove diacritics for matching purposes
                simplified = base
                for char, replacement in [
                    ('â', 'a'), ('î', 'i'), ('û', 'u'), ('ô', 'o'),
                    ('ä', 'a'), ('ë', 'e'), ('ï', 'i'), ('ö', 'o'), ('ü', 'u'),
                    ('é', 'e'), ('è', 'e'), ('ê', 'e'), ('à', 'a'),
                    ('â', 'a'), ('î', 'i'), ('ô', 'o'), ('û', 'u'),
                ]:
                    simplified = simplified.replace(char, replacement)
                return simplified
            
            # Step 1: Build occurrence tracking for original words
            # Maps normalized_word -> list of original indices where it appears
            orig_occurrences = {}
            for i, word in enumerate(original_words):
                norm_word = normalize_word_for_tracking(word)
                if norm_word not in orig_occurrences:
                    orig_occurrences[norm_word] = []
                orig_occurrences[norm_word].append(i)
            
            # Get character-level alignment first (used in Steps 2 and 3)
            norm_to_orig_word = self._align_normalized_to_original_charlevel(
                full_text, normalized_full
            )
            
            # Step 2: Map each normalized word to original
            # Use alignment as primary method. For word splits (multiple normalized
            # words mapping to the same original), the split detection in Step 3
            # handles keeping them on the same line.
            #
            # The occurrence-based approach was removed because it caused incorrect
            # mappings: when a word like "daz" appears multiple times in the original,
            # counting global occurrences led to mapping normalized words to wrong
            # positions (e.g., line 10's "daz" mapped to line 4's "das").
            #
            # The character-level alignment correctly handles:
            # - Word splits (e.g., "inalten" → "in alten")
            # - Word merges
            # - Position-based mapping for non-split words
            norm_to_orig_occurrence = norm_to_orig_word.copy()
            
            # Step 3: Use alignment to detect word splits
            # A true word split is when consecutive normalized words map to the same original
            # AND the original word is longer (contains the split parts)
            # We filter out false splits caused by repeated words like "dise"
            # (norm_to_orig_word already computed above for fallback)
            
            # Detect split groups: consecutive normalized words that map to the same original
            # These should stay on the same line
            # BUT only if it's a true word split (original word contains the parts)
            split_group = {}  # norm_idx -> group_id
            group_id = 0
            i = 0
            while i < total_norm:
                if i == total_norm - 1:
                    # Last word, no split possible
                    split_group[i] = group_id
                    break
                
                # Check if current and next word map to same original (potential word split)
                if norm_to_orig_word[i] == norm_to_orig_word[i + 1]:
                    current_orig = norm_to_orig_word[i]
                    orig_word = original_words[current_orig]
                    
                    # Collect consecutive words mapping to this original
                    split_words = []
                    j = i
                    while j < total_norm and norm_to_orig_word[j] == current_orig:
                        split_words.append(normalized_words[j])
                        j += 1
                    
                    # Check if this is a true word split:
                    # The combined normalized words should be similar to the original
                    # (allowing for minor spelling differences)
                    combined = ''.join(split_words)
                    is_true_split = (
                        len(combined) >= len(orig_word) * 0.5 and  # Combined is at least half the original length
                        len(combined) <= len(orig_word) * 1.5 and  # Combined is at most 1.5x original length
                        len(split_words) > 1  # At least 2 parts
                    )
                    
                    if is_true_split:
                        # True word split - keep parts together
                        for k in range(i, j):
                            split_group[k] = group_id
                        group_id += 1
                        i = j
                    else:
                        # False split (repeated words) - assign each word its own group
                        for k in range(i, j):
                            split_group[k] = group_id
                            group_id += 1
                        i = j
                else:
                    split_group[i] = group_id
                    group_id += 1
                    i += 1
            
            
            # Assign words to lines using word-level alignment
            norm_to_line = self._assign_lines_by_word_alignment(
                lowered_lines if lowered_lines is not None else
                (stripped_lines if stripped_lines is not None else texts),
                normalized_full
            )

            line_word_lists = [[] for _ in range(num_lines)]
            for norm_idx in range(total_norm):
                target_line = norm_to_line[norm_idx]
                line_word_lists[target_line].append(normalized_words[norm_idx])

            # Ensure no empty lines by redistributing words if necessary
            # This can happen when the model produces very different structure
            empty_lines = [i for i in range(num_lines) if not line_word_lists[i]]
            if empty_lines:
                # Find lines with multiple words and move some to empty lines
                for empty_idx in empty_lines:
                    # Find the nearest line with words
                    best_source = None
                    best_dist = float('inf')
                    for src_idx, words in enumerate(line_word_lists):
                        if len(words) > 1:
                            dist = abs(src_idx - empty_idx)
                            if dist < best_dist:
                                best_dist = dist
                                best_source = src_idx
                    
                    if best_source is not None:
                        # Move the last word from source to empty line
                        word = line_word_lists[best_source].pop()
                        line_word_lists[empty_idx].append(word)
            
            normalized = []
            for line_idx in range(num_lines):
                line_text = ' '.join(line_word_lists[line_idx])
                
                # Restore capitalization first (normalized output is lowercase)
                if preserve_capitalization and line_cap_maps is not None and line_idx < len(line_cap_maps):
                    cap_ref = lowered_lines[line_idx] if lowered_lines is not None else None
                    if cap_ref is not None:
                        line_text = self._restore_capitalization_aligned(
                            cap_ref, line_text, line_cap_maps[line_idx]
                        )
                
                # Then restore punctuation (on capitalized text)
                if preserve_punctuation and line_punct_maps is not None and line_idx < len(line_punct_maps):
                    punct_ref = stripped_lines[line_idx] if stripped_lines is not None else None
                    if punct_ref is not None:
                        line_text = self._restore_punctuation_aligned(
                            punct_ref, line_text, line_punct_maps[line_idx]
                        )
                
                normalized.append(line_text)
            
            # Ensure we have the right number of output lines
            while len(normalized) < total_texts:
                normalized.append('')
        else:
            # Original behavior: normalize each line independently
            # Forward preserve_punctuation and preserve_capitalization explicitly
            # so normalize_text doesn't fall back to the config default
            kwargs['preserve_punctuation'] = preserve_punctuation
            kwargs['preserve_capitalization'] = preserve_capitalization
            print()
            normalized = self.normalize_texts(
                texts,
                batch_size=batch_size,
                show_progress=True,
                **kwargs
            )
        
        # Reinsert empty lines at their original positions.
        # normalized[] contains results for non-empty lines only; we need
        # to merge them back into the full line sequence, preserving empty
        # lines at the same positions as the input.
        total_original_lines = len(all_lines)
        if empty_line_indices:
            full_output = [''] * total_original_lines
            for norm_idx, orig_idx in enumerate(non_empty_to_original):
                if norm_idx < len(normalized) and orig_idx < total_original_lines:
                    full_output[orig_idx] = normalized[norm_idx]
            # Any remaining normalized lines beyond the mapping go at the end
            for norm_idx in range(len(non_empty_to_original), len(normalized)):
                full_output.append(normalized[norm_idx])
            output_lines = full_output
        else:
            output_lines = normalized
        
        # Write output
        print(f"\nWriting to: {output_path}")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            for text in output_lines:
                f.write(text + '\n')
        
        print("Done!")
    
    def normalize_multiline_text(
        self,
        input_text: str,
        preserve_punctuation: Optional[bool] = None,
        preserve_capitalization: Optional[bool] = None,
        attach_en_proclitic: Optional[bool] = None,
        lenition_t_after_n: Optional[bool] = None,
        lenition_t_after_l: Optional[bool] = None,
        niet_to_niht: Optional[bool] = None,
        common_apocopes: Optional[bool] = None,
        progress_callback: Optional[callable] = None,
        **kwargs
    ) -> str:
        """
        Normalize multi-line text preserving line structure with cross-line context.
        Replicates the across_lines=True logic from normalize_file() but operates
        on a string rather than a file.

        Returns normalized text with the same number of lines as input.
        """
        if self.config is not None:
            if preserve_punctuation is None:
                preserve_punctuation = self.config.get('preserve_editorial_punctuation', False)
            if preserve_capitalization is None:
                preserve_capitalization = self.config.get('preserve_capitalization', False)
            if attach_en_proclitic is None:
                attach_en_proclitic = self.config.get('post_processing', {}).get('attach_en_proclitic', False)
            if lenition_t_after_n is None:
                lenition_t_after_n = self.config.get('post_processing', {}).get('lenition_t_after_n', False)
            if lenition_t_after_l is None:
                lenition_t_after_l = self.config.get('post_processing', {}).get('lenition_t_after_l', False)
            if niet_to_niht is None:
                niet_to_niht = self.config.get('post_processing', {}).get('niet_to_niht', False)
            if common_apocopes is None:
                common_apocopes = self.config.get('post_processing', {}).get('common_apocopes', False)
        else:
            preserve_punctuation = preserve_punctuation if preserve_punctuation is not None else False
            preserve_capitalization = preserve_capitalization if preserve_capitalization is not None else False
            attach_en_proclitic = attach_en_proclitic if attach_en_proclitic is not None else False
            lenition_t_after_n = lenition_t_after_n if lenition_t_after_n is not None else False
            lenition_t_after_l = lenition_t_after_l if lenition_t_after_l is not None else False
            niet_to_niht = niet_to_niht if niet_to_niht is not None else False
            common_apocopes = common_apocopes if common_apocopes is not None else False

        all_lines = [line.rstrip('\n').rstrip('\r') for line in input_text.split('\n')]

        # Build empty-line map and extract non-empty lines
        empty_line_indices = set()
        texts = []
        for i, line in enumerate(all_lines):
            if line.strip() == '':
                empty_line_indices.add(i)
            else:
                texts.append(line.strip())

        non_empty_to_original = []
        for i, line in enumerate(all_lines):
            if line.strip() != '':
                non_empty_to_original.append(i)

        if not preserve_capitalization:
            texts = [line.lower() for line in texts]

        total_texts = len(texts)
        if total_texts == 0:
            return '\n'.join(all_lines)

        # across_lines=True logic
        line_punct_maps = None
        line_cap_maps = None
        stripped_lines = None
        lowered_lines = None

        if preserve_punctuation:
            line_punct_maps = []
            stripped_lines = []
            for line in texts:
                stripped_line, punct_map = self._extract_punctuation(line)
                stripped_lines.append(stripped_line)
                line_punct_maps.append(punct_map)

        if preserve_capitalization:
            line_cap_maps = []
            lowered_lines = []
            cap_source_lines = stripped_lines if stripped_lines is not None else texts
            for line in cap_source_lines:
                lowered_line, cap_map = self._extract_capitalization(line)
                lowered_lines.append(lowered_line)
                line_cap_maps.append(cap_map)

        if lowered_lines is not None:
            word_counts = [len(line.split()) for line in lowered_lines]
            full_text = ' '.join(lowered_lines)
        elif stripped_lines is not None:
            word_counts = [len(line.split()) for line in stripped_lines]
            full_text = ' '.join(stripped_lines)
        else:
            word_counts = [len(line.split()) for line in texts]
            full_text = ' '.join(texts)

        normalized_full = self.normalize_text(
            full_text, preserve_punctuation=False,
            preserve_capitalization=False,
            progress_callback=progress_callback,
            common_apocopes=False,  # applied below on final result
            **kwargs
        )

        num_lines = len(texts)
        normalized_words = normalized_full.split()
        total_norm = len(normalized_words)

        word_to_line = []
        for line_idx, count in enumerate(word_counts):
            for _ in range(count):
                word_to_line.append(line_idx)

        total_orig = sum(word_counts)
        original_words = full_text.split()

        def normalize_word_for_tracking(word: str) -> str:
            base = word.lower().replace('ſ', 's')
            simplified = base
            for char, replacement in [
                ('â', 'a'), ('î', 'i'), ('û', 'u'), ('ô', 'o'),
                ('ä', 'a'), ('ë', 'e'), ('ï', 'i'), ('ö', 'o'), ('ü', 'u'),
                ('é', 'e'), ('è', 'e'), ('ê', 'e'), ('à', 'a'),
            ]:
                simplified = simplified.replace(char, replacement)
            return simplified

        orig_occurrences = {}
        for i, word in enumerate(original_words):
            norm_word = normalize_word_for_tracking(word)
            if norm_word not in orig_occurrences:
                orig_occurrences[norm_word] = []
            orig_occurrences[norm_word].append(i)

        # Progress: alignment is the slowest step on CPU (O(n*m) difflib)
        if progress_callback:
            progress_callback(78, 100)

        # Progress: about to start character-level alignment (slowest CPU step)
        if progress_callback:
            progress_callback(79, 100)

        norm_to_orig_word = self._align_normalized_to_original_charlevel(
            full_text, normalized_full
        )

        # Progress: alignment done
        if progress_callback:
            progress_callback(84, 100)

        norm_to_orig_occurrence = norm_to_orig_word.copy()

        split_group = {}
        group_id = 0
        i = 0
        while i < total_norm:
            if i == total_norm - 1:
                split_group[i] = group_id
                group_id += 1
                break

            # Only group consecutive normalized words that map to the *same*
            # original word (true word splits like "inalten" → "in alten").
            # Do NOT group words mapping to adjacent originals — that collapses
            # normal line boundaries and piles words onto the first line.
            orig_i = norm_to_orig_word[i]
            orig_next = norm_to_orig_word[i + 1]

            if orig_i == orig_next:
                # Collect consecutive words mapping to this original
                split_words = []
                j = i
                while j < total_norm and norm_to_orig_word[j] == orig_i:
                    split_words.append(normalized_words[j])
                    j += 1

                # Validate: true splits have combined length similar to original
                combined = ''.join(split_words)
                orig_word = original_words[orig_i] if orig_i < len(original_words) else ''
                is_true_split = (
                    len(split_words) > 1 and
                    len(combined) >= len(orig_word) * 0.5 and
                    len(combined) <= len(orig_word) * 1.5
                )

                if is_true_split:
                    for k in range(i, j):
                        split_group[k] = group_id
                    group_id += 1
                    i = j
                else:
                    # False split (repeated words) — each gets its own group
                    split_group[i] = group_id
                    group_id += 1
                    i += 1
            else:
                split_group[i] = group_id
                group_id += 1
                i += 1

        # Progress: split grouping done, starting line assignment
        if progress_callback:
            progress_callback(85, 100)

        norm_to_line = self._assign_lines_by_word_alignment(
            lowered_lines if lowered_lines is not None else
            (stripped_lines if stripped_lines is not None else texts),
            normalized_full
        )

        # Progress: line assignment done
        if progress_callback:
            progress_callback(88, 100)

        line_word_lists = [[] for _ in range(num_lines)]
        for norm_idx in range(total_norm):
            target_line = norm_to_line[norm_idx]
            line_word_lists[target_line].append(normalized_words[norm_idx])

        # Progress: word distribution done
        if progress_callback:
            progress_callback(90, 100)

        empty_lines = [i for i in range(num_lines) if not line_word_lists[i]]
        if empty_lines:
            for empty_idx in empty_lines:
                best_source = None
                best_dist = float('inf')
                for src_idx, words in enumerate(line_word_lists):
                    if len(words) > 1:
                        dist = abs(src_idx - empty_idx)
                        if dist < best_dist:
                            best_dist = dist
                            best_source = src_idx

                if best_source is not None:
                    word = line_word_lists[best_source].pop()
                    line_word_lists[empty_idx].append(word)

        # Progress: starting capitalization/punctuation restoration loop
        if progress_callback:
            progress_callback(91, 100)

        normalized = []
        for line_idx in range(num_lines):
            line_text = ' '.join(line_word_lists[line_idx])

            if preserve_capitalization and line_cap_maps is not None and line_idx < len(line_cap_maps):
                cap_ref = lowered_lines[line_idx] if lowered_lines is not None else None
                if cap_ref is not None:
                    line_text = self._restore_capitalization_aligned(
                        cap_ref, line_text, line_cap_maps[line_idx]
                    )

            if preserve_punctuation and line_punct_maps is not None and line_idx < len(line_punct_maps):
                punct_ref = stripped_lines[line_idx] if stripped_lines is not None else None
                if punct_ref is not None:
                    line_text = self._restore_punctuation_aligned(
                        punct_ref, line_text, line_punct_maps[line_idx]
                    )

            normalized.append(line_text)

        # Progress: restoration loop done
        if progress_callback:
            progress_callback(93, 100)

        while len(normalized) < total_texts:
            normalized.append('')

        # Reinsert empty lines
        total_original_lines = len(all_lines)
        if empty_line_indices:
            full_output = [''] * total_original_lines
            for norm_idx, orig_idx in enumerate(non_empty_to_original):
                if norm_idx < len(normalized) and orig_idx < total_original_lines:
                    full_output[orig_idx] = normalized[norm_idx]
            for norm_idx in range(len(non_empty_to_original), len(normalized)):
                full_output.append(normalized[norm_idx])
            output_lines = full_output
        else:
            output_lines = normalized

        # Progress: empty line reinsertion done
        if progress_callback:
            progress_callback(95, 100)

        result = '\n'.join(output_lines)

        # Progress: join done, starting post-processing regex
        if progress_callback:
            progress_callback(96, 100)

        # Post-processing: common apocopes (before en-attachment so "en wile" → "en wil" → "enwil")
        if common_apocopes:
            result = _common_apocopes(result)
        # Post-processing: attach negation proclitic (after apocopes)
        if attach_en_proclitic:
            result = _attach_en_proclitic(result)
        # Post-processing: lenition of t after n before vowel
        if lenition_t_after_n:
            result = _lenition_t_after_n_before_vowel(result)
        # Post-processing: lenition of t after l before vowel
        if lenition_t_after_l:
            result = _lenition_t_after_l_before_vowel(result)
        # Post-processing: niet → niht
        if niet_to_niht:
            result = _niet_to_niht(result)

        # Progress: all done
        if progress_callback:
            progress_callback(100, 100)

        return result

    def normalize_and_compare(
        self,
        texts: List[str],
        reference_texts: Optional[List[str]] = None,
        num_samples: int = 5,
        **kwargs
    ):
        """
        Normalize texts and optionally compare with references.

        Args:
            texts: List of input texts
            reference_texts: Optional list of reference normalized texts
            num_samples: Number of samples to display
            **kwargs: Additional arguments for normalize_text
        """
        print(f"\n{'='*60}")
        print(f"Normalizing {len(texts)} texts")
        print(f"{'='*60}\n")
        
        normalized = self.normalize_texts(texts, **kwargs)
        
        # Print samples
        num_samples = min(num_samples, len(texts))
        
        for i in range(num_samples):
            print(f"Sample {i + 1}:")
            print(f"  Original:    {texts[i]}")
            print(f"  Normalized:  {normalized[i]}")
            
            if reference_texts:
                print(f"  Reference:   {reference_texts[i]}")
                match = normalized[i] == reference_texts[i]
                print(f"  Match:       {match}")
            
            print()
        
        return normalized


def main():
    parser = argparse.ArgumentParser(description='Normalize Middle High German text')
    
    # Model arguments
    parser.add_argument('--model', type=str,
                       default='/home/jonas/normaere/checkpoints/best_model',
                       help='Path to trained model')
    parser.add_argument('--config', type=str, default='/home/jonas/normaere/config.yaml',
                       help='Path to configuration file')
    
    # Input arguments
    parser.add_argument('--text', type=str, help='Single text to normalize')
    parser.add_argument('--input_file', type=str, help='Input file with texts (one per line)')
    parser.add_argument('--output_file', type=str, help='Output file for normalized texts')
    
    # Generation arguments
    parser.add_argument('--max_length', type=int, default=512,
                       help='Maximum generation length')
    parser.add_argument('--num_beams', type=int, default=5,
                       help='Number of beams')
    parser.add_argument('--length_penalty', type=float, default=0.6,
                       help='Length penalty')
    parser.add_argument('--do_sample', action='store_true',
                       help='Use sampling instead of beam search')
    parser.add_argument('--top_p', type=float, default=None,
                       help='Nucleus sampling threshold (0.0 to 1.0)')
    parser.add_argument('--temperature', type=float, default=None,
                       help='Temperature for sampling (0.1 to 2.0, lower = more deterministic, higher = more random)')
    parser.add_argument('--batch_size', type=int, default=8,
                       help='Batch size for file processing')
    
    # Stride arguments
    parser.add_argument('--stride', action='store_true', default=True,
                       help='Use sliding window stride for long texts (default: True)')
    parser.add_argument('--no-stride', action='store_false', dest='stride',
                       help='Disable stride; truncate long texts instead')
    parser.add_argument('--stride_window_tokens', type=int, default=None,
                       help='Max input tokens per stride window (default: 400)')
    parser.add_argument('--stride_step_tokens', type=int, default=None,
                       help='Token advance between windows; overlap = window - step (default: 300)')
    
    # Across-lines arguments
    parser.add_argument('--across-lines', action='store_true', default=True,
                       help='Join all lines into one continuous text for normalization, '
                            'utilizing full model capacity and stride across line boundaries '
                            '(default: True)')
    parser.add_argument('--no-across-lines', action='store_false', dest='across_lines',
                       help='Normalize each line independently (original behavior)')
    
    # Punctuation preservation arguments
    parser.add_argument('--preserve-punctuation', action='store_true', default=True,
                       help='Strip editorial punctuation before normalization and '
                            'restore it afterwards (default: True)')
    parser.add_argument('--no-preserve-punctuation', action='store_false', dest='preserve_punctuation',
                       help='Do not strip/restore punctuation; pass text as-is to the model')
    
    # Capitalization preservation arguments
    parser.add_argument('--preserve-capitalization', action='store_true', default=True,
                       help='Record capitalization before normalization (lowercasing input) '
                            'and restore it afterwards (default: True)')
    parser.add_argument('--no-preserve-capitalization', action='store_false', dest='preserve_capitalization',
                       help='Do not preserve capitalization; input is lowercased without restoration')
    
    args = parser.parse_args()
    
    # Initialize normalizer
    normalizer = MHGNormalizer(args.model, args.config)
    
    # Build stride kwargs
    stride_kwargs = {
        'stride': args.stride,
        'stride_window_tokens': args.stride_window_tokens,
        'stride_step_tokens': args.stride_step_tokens,
    }
    
    if args.text:
        # Normalize single text
        print("\nOriginal text:")
        print(args.text)
        print("\nNormalized text:")
        normalized = normalizer.normalize_text(
            args.text,
            max_length=args.max_length,
            num_beams=args.num_beams,
            length_penalty=args.length_penalty,
            do_sample=args.do_sample,
            top_p=args.top_p,
            temperature=args.temperature,
            preserve_punctuation=args.preserve_punctuation,
            preserve_capitalization=args.preserve_capitalization,
            **stride_kwargs
        )
        print(normalized)
        print()
    
    elif args.input_file:
        # Normalize file
        if not args.output_file:
            # Create output path based on input
            input_path = Path(args.input_file)
            args.output_file = input_path.parent / f"{input_path.stem}_normalized{input_path.suffix}"
        
        normalizer.normalize_file(
            args.input_file,
            args.output_file,
            batch_size=args.batch_size,
            across_lines=args.across_lines,
            preserve_punctuation=args.preserve_punctuation,
            preserve_capitalization=args.preserve_capitalization,
            max_length=args.max_length,
            num_beams=args.num_beams,
            length_penalty=args.length_penalty,
            do_sample=args.do_sample,
            top_p=args.top_p,
            temperature=args.temperature,
            **stride_kwargs
        )
    else:
        # Interactive mode
        print("\nInteractive mode. Enter MHG text to normalize (Ctrl+D to exit):")
        print("="*60 + "\n")
        
        try:
            texts = []
            while True:
                text = input("> ")
                if text.strip():
                    texts.append(text)
        except EOFError:
            pass
        
        if texts:
            print(f"\nNormalizing {len(texts)} texts...\n")
            normalized = normalizer.normalize_texts(
                texts,
                preserve_punctuation=args.preserve_punctuation,
                preserve_capitalization=args.preserve_capitalization,
                **stride_kwargs
            )
            
            for i, (orig, norm) in enumerate(zip(texts, normalized)):
                print(f"\nText {i + 1}:")
                print(f"  Original:   {orig}")
                print(f"  Normalized: {norm}")


if __name__ == '__main__':
    main()