import gc
import os
import platform
import struct
import pickle
import json
import datetime
import time
import logging
import unicodedata
import importlib.resources
import regex
import re
import multiprocessing

from .params import (
    METASPACE,
    SPECIAL_TOKENS as _SPECIAL_TOKENS_TUPLE,
    PAD_TOKEN,
    PAD_TOKEN_ID,
    REGEX_PATTERN_METASPACE,
    REGEX_PATTERN_GPT2,
    DEFAULT_VOCAB_SIZE,
    MIN_VOCAB_SIZE,
    DEFAULT_MIN_FREQUENCY,
    DEFAULT_SPACE_STRATEGY,
    SPACE_STRATEGY_METASPACE,
    SPACE_STRATEGY_GPT2,
    KINYARWANDA_PREFIXES,
    LOAN_WORD_ALLOWLIST,
    LARGE_CORPUS_CPU_CAP,
    RUST_MERGE_BATCH_SIZE,
    CHECKPOINT_EVERY_N_MERGES,
    PYTHON_FALLBACK_CHECKPOINT_EVERY,
    DECODE_MULTIPROCESS_THRESHOLD,
)

try:
    from kin_merge import (
        rust_bpe_train_loop,
        rust_create_tokens_parallel,
        rust_encode_batch,
        rust_encode_text,
        rust_encode_texts,
        rust_pretokenize_text,
        rust_preprocess_text,
        rust_filter_words,
    )
    _HAS_RUST_MERGER = True
except Exception:
    rust_bpe_train_loop = None
    rust_create_tokens_parallel = None
    rust_encode_batch = None
    rust_encode_text = None
    rust_encode_texts = None
    rust_pretokenize_text = None
    rust_preprocess_text = None
    rust_filter_words = None
    _HAS_RUST_MERGER = False

logger = logging.getLogger(__name__)

# Re-export for backward compatibility — canonical definitions live in params.py.
_SPECIAL_TOKENS = _SPECIAL_TOKENS_TUPLE

# Pre-compiled regex patterns — selected based on space_strategy at load/init time.
# Updated from tiktoken o200k_base insights: \p{M} captures combining marks
# (tone accents, diacritics) so accented Kinyarwanda stays intact; \p{N}+
# replaces {1,3} to avoid splitting multi-digit numbers.
_COMPILED_PATTERN_METASPACE = regex.compile(REGEX_PATTERN_METASPACE, regex.UNICODE)
_COMPILED_PATTERN_GPT2 = regex.compile(REGEX_PATTERN_GPT2, regex.UNICODE)

# ───────────────────────────── Safe Token Chunk I/O ────────────────────────
# Replaces pickle for current_token_chunks storage.  Format:
#   4 bytes  magic  b"KINC"
#   4 bytes  version (u32 LE) — currently 1
#   8 bytes  number of chunks (u64 LE)
#   For each chunk:
#       4 bytes  chunk length L (u32 LE)
#       L*4 bytes  token values (each u32 LE)
#
# Same compactness as pickle (~4 bytes/token) with zero code-execution risk.
# ───────────────────────────────────────────────────────────────────────────────

_TOKEN_CHUNKS_MAGIC = b"KINC"
_TOKEN_CHUNKS_VERSION = 1
_TOKEN_CHUNKS_EXT = ".bin"
_TOKEN_CHUNKS_LEGACY_EXT = ".pkl"


def _save_token_chunks(chunks, path):
    """Write token chunks to a safe binary file (no pickle)."""
    with open(path, "wb") as f:
        f.write(_TOKEN_CHUNKS_MAGIC)
        f.write(struct.pack("<I", _TOKEN_CHUNKS_VERSION))
        f.write(struct.pack("<Q", len(chunks)))
        for chunk in chunks:
            f.write(struct.pack("<I", len(chunk)))
            f.write(struct.pack(f"<{len(chunk)}I", *chunk))


def _load_token_chunks(path):
    """Read token chunks from the safe binary format."""
    with open(path, "rb") as f:
        magic = f.read(4)
        if magic != _TOKEN_CHUNKS_MAGIC:
            raise ValueError(
                f"Invalid token chunk file (bad magic): {path}. "
                "Expected safe binary format (.bin). If you have a legacy "
                ".pkl file, re-run first-train to generate a new checkpoint."
            )
        (version,) = struct.unpack("<I", f.read(4))
        if version != _TOKEN_CHUNKS_VERSION:
            raise ValueError(f"Unsupported token chunk version {version} in {path}")
        (num_chunks,) = struct.unpack("<Q", f.read(8))
        chunks = []
        for _ in range(num_chunks):
            (length,) = struct.unpack("<I", f.read(4))
            data = struct.unpack(f"<{length}I", f.read(length * 4))
            chunks.append(list(data))
        return chunks


def _token_chunks_path(directory):
    """Return the canonical safe-binary chunk path for a checkpoint directory."""
    return os.path.join(directory, "current_token_chunks" + _TOKEN_CHUNKS_EXT)


"""
Kin_tokenizer is a python module which includes class and methods
for Kinyarwanda tokenizer
"""


def _create_tokens_standalone(text):
    """
    Standalone function for multiprocessing to create tokens from text.
    Does NOT instantiate KinTokenizer — avoids expensive pkl reload per worker.
    """
    if isinstance(text, str):
        text = text.encode("UTF-8")
    elif not isinstance(text, bytes):
        raise ValueError("Expected string or bytes")
    return list(map(int, text))


def _merge_tokens_standalone(args):
    """
    Standalone function for multiprocessing to merge tokens.
    Does NOT instantiate KinTokenizer — pure function on data.
    Args: tuple of (pair, tokens, new_token).
    """
    pair, tokens, new_token = args
    new_tokens = []
    index = 0
    while index < len(tokens):
        if index < len(tokens) - 1 and pair[0] == tokens[index] and pair[1] == tokens[index + 1]:
            new_tokens.append(new_token)
            index += 2
        else:
            new_tokens.append(tokens[index])
            index += 1
    return new_tokens


# Backward compatibility aliases (deprecated — will be removed in v4)
parallel_create_tokens = _create_tokens_standalone
parallel_merge_tokens = _merge_tokens_standalone


# ──────────────────────────────────────────────────────────────────────
# Standalone encode worker ― avoids pickling KinTokenizer on every task.
# The worker is initialised once per Pool with the merged_tokens dict;
# subsequent map() calls need only pickle the (index, word) tuples.
# ──────────────────────────────────────────────────────────────────────

_encode_worker_merged_tokens = None  # module-level cache populated by _init_encode_worker


def _init_encode_worker(merged_tokens: dict):
    """Pool initializer: stash merged_tokens in a per-process global."""
    global _encode_worker_merged_tokens
    _encode_worker_merged_tokens = merged_tokens


def _encode_chunk_standalone(indexed_word):
    """
    Standalone BPE encode of a single pre-tokenised chunk (word/token).
    Uses the process-global _encode_worker_merged_tokens populated by the
    pool initializer, so KinTokenizer is NOT pickled on every task dispatch.

    Args:
        indexed_word: (index: int, word: str) tuple

    Returns:
        (index, list[int]) — index preserved for stable reassembly.
    """
    index, word = indexed_word
    # UTF-8 byte encode
    if isinstance(word, str):
        tokens = list(word.encode('UTF-8'))
    elif isinstance(word, bytes):
        tokens = list(word)
    else:
        raise ValueError(f"Expected str or bytes, got {type(word)}")

    merged = _encode_worker_merged_tokens  # fast local alias
    # FIX #9: Scan for the pair with the lowest merge rank directly,
    # instead of building an intermediate stats dict (pair→count).
    # The count is unused — only the rank (token ID) in merged_tokens matters.
    while len(tokens) > 1:
        best_pair = None
        best_rank = float("inf")
        for i in range(len(tokens) - 1):
            pair = (tokens[i], tokens[i + 1])
            rank = merged.get(pair, float("inf"))
            if rank < best_rank:
                best_rank = rank
                best_pair = pair
        if best_pair is None or best_rank == float("inf"):
            break
        tokens = _merge_tokens_standalone((best_pair, tokens, merged[best_pair]))
    return index, tokens


# ──────────────────────────────────────────────────────────────────────
# Standalone decode worker ― avoids pickling KinTokenizer on every task.
# Same pattern as the encode worker above.
# ──────────────────────────────────────────────────────────────────────

_decode_worker_vocab = None  # module-level cache populated by _init_decode_worker


def _init_decode_worker(vocab: dict):
    """Pool initializer: stash vocab in a per-process global."""
    global _decode_worker_vocab
    _decode_worker_vocab = vocab


def _decode_chunk_standalone(args):
    """Standalone decode of a token chunk using the process-global vocab."""
    chunk, return_eos = args
    vocab = _decode_worker_vocab
    tokens_part = []
    has_eos = False
    invalid_indices = []
    for idx in chunk:
        if idx not in vocab:
            invalid_indices.append(idx)
            continue
        if return_eos and vocab[idx] == "<|EOS|>":
            has_eos = True
            continue
        tokens_part.append(vocab[idx])
    return tokens_part, has_eos, invalid_indices


def _cpu_count():
    count = None
    process_cpu_count = getattr(os, "process_cpu_count", None)
    if callable(process_cpu_count):
        try:
            count = process_cpu_count()
        except Exception:
            count = None
    if count is None:
        count = os.cpu_count()
    return max(1, int(count or 1))


def _get_mp_context():
    """Get the appropriate multiprocessing context for the current platform."""
    system = platform.system()
    if system in ("Windows", "Darwin"):
        return multiprocessing.get_context('spawn')
    try:
        return multiprocessing.get_context('forkserver')
    except ValueError:
        return multiprocessing.get_context('spawn')


def _select_top_pair(stats):
    """
    Deterministically choose the most frequent pair.
    Ties are resolved by lexicographically smaller pair for stable training.
    VERIFIED #4: Matches Rust select_top_pair / select_top_pair_signed which
    uses b.0.cmp(&a.0) (reverse tuple cmp) — both pick smallest pair on tie.
    """
    return max(stats, key=lambda pair: (stats[pair], -pair[0], -pair[1]))


class KinTokenizer:
    def __init__(self, autoload=True):
        """
        Initialize KinTokenizer.
        Attempts to load a pre-trained tokenizer state from the package data.

        Vocabulary Index Design (FIX #14):
        - ID 0:       <|PAD|> — reserved padding token
        - IDs 1-255:  raw byte tokens (populated during train() or load())
        - IDs 256+:   BPE merge tokens, assigned in merge order (lower ID = higher priority merge)
        - Last 6 IDs: special tokens (<|EOS|>, <|BOS|>, <|SEP|>, <|MASK|>, <|UNK|>, <|CLS|>)

        This layout means token IDs double as merge priority ranks in the BPE encoder:
        encode_piece_ids picks the pair whose merged token has the lowest ID.
        """
        super(KinTokenizer, self).__init__()

        self.__vocab = {PAD_TOKEN_ID: PAD_TOKEN}
        self.normalization = {"lowercase": False}
        self.merged_tokens = {}
        self.vocab_size = None
        self._token_to_id_cache = None
        # SPACE STRATEGY: SentencePiece metaspace (Option A).
        # Every word gets a ▁ prefix so BPE learns one canonical form per morpheme.
        self.space_strategy = DEFAULT_SPACE_STRATEGY

        # Kinyarwanda morpheme-aware prefix list from params.
        self.kinyarwanda_prefixes = list(KINYARWANDA_PREFIXES)

        # Select the pre-compiled pattern for the active space strategy.
        # New training uses metaspace; old checkpoints loaded via load() switch to gpt2 compat.
        self.compiled_pattern = _COMPILED_PATTERN_METASPACE

        if autoload:
            # Prefer the repo-level checkpoint when running from source.
            source_tree_checkpoint = os.path.abspath(
                os.path.join(os.path.dirname(__file__), "..", "data", "kin_tokenizer.json")
            )
            if os.path.isfile(source_tree_checkpoint):
                self.load(source_tree_checkpoint, quiet_missing=True)

            # Fall back to the default checkpoint directory (data/checkpoint/).
            if self.vocab_size is None:
                checkpoint_dir = os.path.abspath(
                    os.path.join(os.path.dirname(__file__), "..", "data", "checkpoint", "kin_tokenizer.json")
                )
                if os.path.isfile(checkpoint_dir):
                    self.load(checkpoint_dir, quiet_missing=True)

            # Fall back to packaged JSON artifact.
            if self.vocab_size is None:
                tokenizer_json_path_obj = importlib.resources.files('kin_tokenizer').joinpath('data/kin_tokenizer.json')
                with importlib.resources.as_file(tokenizer_json_path_obj) as tokenizer_json_path:
                    self.load(tokenizer_json_path, quiet_missing=True)

            if self.vocab_size is None and os.getenv("KIN_TOKENIZER_ALLOW_PICKLE", "0") == "1":
                tokenizer_pkl_path_obj = importlib.resources.files('kin_tokenizer').joinpath('data/kin_tokenizer.pkl')
                with importlib.resources.as_file(tokenizer_pkl_path_obj) as tokenizer_pkl_path:
                    self.load(tokenizer_pkl_path, allow_pickle=True, quiet_missing=True)

    @property
    def vocab(self):
        return {
            key: value.decode("UTF-8", errors="replace") if isinstance(value, bytes) else value
            for key, value in self.__vocab.items()
        }

    def set_vocab(self, vocab):
        """
        Set the vocabulary of the tokenizer.
        vocab: dictionary of {int: bytes}
        """
        if len(self.__vocab) <= 1:
            if isinstance(vocab, dict):
                self.__vocab = vocab
                self._reset_token_lookup_cache()
            else:
                raise ValueError("Expected a dictionary of {integer: bytes}")
        else:
            raise ValueError("Vocab cannot be overridden")

    # Removed deprecated alias set_vacab in v4; use set_vocab directly.

    def set_merged_tokens(self, merged_tokens):
        """
        Set merged_tokens.
        merged_tokens: dictionary of {(int, int): int}
        """
        if len(self.merged_tokens) < 1:
            if isinstance(merged_tokens, dict):
                self.merged_tokens = merged_tokens
            else:
                raise ValueError("Expected a dictionary of {(integer, integer): integer}")
        else:
            raise ValueError("merged_tokens cannot be overridden")

    def _reset_token_lookup_cache(self):
        self._token_to_id_cache = None

    def reset_for_training(self):
        """Reset to a clean fresh-training state with only PAD in vocab."""
        self.__vocab = {PAD_TOKEN_ID: PAD_TOKEN}
        self.normalization = {"lowercase": False}
        self.merged_tokens = {}
        self.vocab_size = None
        self.space_strategy = DEFAULT_SPACE_STRATEGY
        self.compiled_pattern = _COMPILED_PATTERN_METASPACE
        self._reset_token_lookup_cache()

    def ensure_base_byte_vocab(self):
        """Ensure byte fallback tokens 1..255 exist exactly once."""
        for index in range(1, 256):
            self.__vocab.setdefault(index, bytes([index]))

    def ensure_special_tokens(self):
        """Add missing special tokens without duplicating existing ones."""
        existing = {
            value for value in self.__vocab.values()
            if isinstance(value, str)
        }
        max_token = max(self.__vocab.keys(), default=0)
        for token in _SPECIAL_TOKENS:
            if token in existing:
                continue
            max_token += 1
            self.__vocab[max_token] = token
            existing.add(token)

    def _should_lowercase(self, lowercase=None):
        if lowercase is None:
            return bool(self.normalization.get("lowercase", False))
        return bool(lowercase)

    def _normalize_text_for_encoding(self, text, lowercase=None, _skip_nfc=False):
        """Normalize text before encoding to match the distribution seen during training.

        Applies Unicode NFC normalization (canonical decomposition + composition)
        so that equivalent codepoint sequences produce identical byte sequences.
        This is critical for deterministic tokenization — without NFC, the same
        visible character can produce different token IDs depending on whether
        the input uses precomposed or decomposed Unicode forms.

        Optionally applies case-folding based on the tokenizer's normalization config.

        _skip_nfc: when True, skips the NFC pass (caller already guaranteed NFC text,
        e.g. after preprocess_text which does NFC internally).
        """
        if not isinstance(text, str):
            raise ValueError("Expected a string!")
        if not _skip_nfc:
            text = unicodedata.normalize("NFC", text)
        if self._should_lowercase(lowercase):
            text = text.lower()
        # METASPACE (Option A, only when this tokenizer was trained with metaspace):
        # prepend ▁ and replace spaces with ▁ so every word — including the first —
        # gets a consistent ▁ prefix before BPE tokenisation.
        if getattr(self, "space_strategy", DEFAULT_SPACE_STRATEGY) == SPACE_STRATEGY_METASPACE:
            text = (" " + text.lstrip(" ")).replace(" ", METASPACE)
        return text

    def token_to_id(self, token):
        if self._token_to_id_cache is None:
            cache = {}
            for token_id, value in self.__vocab.items():
                if isinstance(value, bytes):
                    decoded = value.decode("UTF-8", errors="replace")
                else:
                    decoded = value
                cache.setdefault(decoded, token_id)
            self._token_to_id_cache = cache
        return self._token_to_id_cache.get(token)

    def id_to_token(self, token_id):
        value = self.__vocab.get(token_id)
        if value is None:
            return None
        if isinstance(value, bytes):
            return value.decode("UTF-8", errors="replace")
        return value

    def _state_as_json(self):
        return {
            "vocab_size": self.vocab_size if self.vocab_size is not None else len(self.__vocab),
            "normalization": {"lowercase": bool(self.normalization.get("lowercase", False))},
            "space_strategy": getattr(self, "space_strategy", DEFAULT_SPACE_STRATEGY),
            "vocab": {str(k): list(v) if isinstance(v, bytes) else v for k, v in self.__vocab.items()},
            "merged_tokens": {f"{k[0]},{k[1]}": v for k, v in self.merged_tokens.items()},
        }

    def _resolve_tokenizer_path(self, tokenizer_path):
        tokenizer_path = os.fspath(tokenizer_path)
        if os.path.isdir(tokenizer_path):
            json_path = os.path.join(tokenizer_path, "kin_tokenizer.json")
            if os.path.exists(json_path):
                return json_path
            pkl_path = os.path.join(tokenizer_path, "kin_tokenizer.pkl")
            if os.path.exists(pkl_path):
                return pkl_path
            return json_path
        return tokenizer_path

    def save(self, path, save_legacy_pickle=False):
        """
        Save the tokenizer state.
        path: directory path where the tokenizer will be saved
        save_legacy_pickle: if True, also writes insecure pickle checkpoint for backward compatibility
        """
        if not os.path.exists(path) or not os.path.isdir(path):
            raise ValueError("The path should be a directory path and it should exist!")

        json_path = os.path.join(path, "kin_tokenizer.json")
        json_state = self._state_as_json()
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(json_state, f, ensure_ascii=False, indent=2)

        if save_legacy_pickle:
            pkl_path = os.path.join(path, "kin_tokenizer.pkl")
            with open(pkl_path, "wb") as f:
                pickle.dump(self, f)

        display_path = json_path.replace("\\", "\\\\")
        print(f"\n\n{'=' * (len(display_path) + 33)}\nTokenizer saved successfully at {display_path}\n{'=' * (len(display_path) + 33)}\n")

    def export_huggingface(self, output_dir):
        """
        Export tokenizer in HuggingFace-compatible format.

        Writes vocab.json, merges.txt, tokenizer_config.json, and
        special_tokens_map.json to output_dir, suitable for loading with
        HuggingFace transformers PreTrainedTokenizerFast or equivalent.

        Args:
            output_dir: directory path for the exported files
        """
        os.makedirs(output_dir, exist_ok=True)

        # Build token_string -> id mapping for vocab.json
        vocab_json = {}
        for tid, value in self.__vocab.items():
            if isinstance(value, bytes):
                # Represent byte tokens using the GPT-2/tiktoken byte encoding convention
                token_str = value.decode("utf-8", errors="replace")
            elif isinstance(value, str):
                token_str = value
            else:
                token_str = str(value)
            vocab_json[token_str] = tid

        with open(os.path.join(output_dir, "vocab.json"), "w", encoding="utf-8") as f:
            json.dump(vocab_json, f, ensure_ascii=False, indent=2)

        # Build merges.txt: merge rules in priority order
        # merged_tokens is {(left_id, right_id): new_id} — sort by new_id for merge priority
        def _id_to_token_str(tid):
            val = self.__vocab.get(tid)
            if val is None:
                return f"<id_{tid}>"
            if isinstance(val, bytes):
                return val.decode("utf-8", errors="replace")
            return str(val)

        sorted_merges = sorted(self.merged_tokens.items(), key=lambda x: x[1])
        with open(os.path.join(output_dir, "merges.txt"), "w", encoding="utf-8") as f:
            f.write("#version: 0.2\n")
            for (left_id, right_id), _ in sorted_merges:
                left_str = _id_to_token_str(left_id)
                right_str = _id_to_token_str(right_id)
                f.write(f"{left_str} {right_str}\n")

        # Identify special token IDs
        special_tokens = {}
        for tid, value in self.__vocab.items():
            if isinstance(value, str) and value.startswith("<|") and value.endswith("|>"):
                name = value[2:-2].lower()  # e.g. "PAD", "EOS" -> "pad", "eos"
                special_tokens[name] = {"id": tid, "content": value}

        # tokenizer_config.json
        tokenizer_config = {
            "model_type": "bpe",
            "tokenizer_class": "PreTrainedTokenizerFast",
            "vocab_size": self.vocab_size or len(self.__vocab),
            "do_lower_case": bool(self.normalization.get("lowercase", False)),
            "pad_token": special_tokens.get("pad", {}).get("content"),
            "eos_token": special_tokens.get("eos", {}).get("content"),
            "bos_token": special_tokens.get("bos", {}).get("content"),
            "unk_token": special_tokens.get("unk", {}).get("content"),
            "sep_token": special_tokens.get("sep", {}).get("content"),
            "cls_token": special_tokens.get("cls", {}).get("content"),
            "mask_token": special_tokens.get("mask", {}).get("content"),
        }
        # Remove None entries
        tokenizer_config = {k: v for k, v in tokenizer_config.items() if v is not None}

        with open(os.path.join(output_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
            json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)

        # special_tokens_map.json
        special_tokens_map = {}
        for name in ("pad", "eos", "bos", "unk", "sep", "cls", "mask"):
            info = special_tokens.get(name)
            if info:
                special_tokens_map[f"{name}_token"] = info["content"]

        with open(os.path.join(output_dir, "special_tokens_map.json"), "w", encoding="utf-8") as f:
            json.dump(special_tokens_map, f, ensure_ascii=False, indent=2)

        print(f"HuggingFace tokenizer exported to: {output_dir}")
        print(f"  vocab.json: {len(vocab_json)} entries")
        print(f"  merges.txt: {len(sorted_merges)} merge rules")
        print(f"  tokenizer_config.json + special_tokens_map.json")

    def load(self, tokenizer_path, allow_pickle=False, quiet_missing=False):
        """
        Load tokenizer state from a directory, JSON file, or legacy pickle file.
        tokenizer_path: directory path or full file path.
        allow_pickle: when True, allows loading legacy pickle checkpoints.
        quiet_missing: suppress missing-file warnings (used during optional auto-load).
        """
        resolved_path = self._resolve_tokenizer_path(tokenizer_path)
        try:
            extension = os.path.splitext(resolved_path)[1].lower()
            if extension == ".json":
                with open(resolved_path, "r", encoding="utf-8") as f:
                    state = json.load(f)

                raw_vocab = state.get("vocab", {})
                vocab = {}
                for key, value in raw_vocab.items():
                    token_id = int(key)
                    if isinstance(value, list):
                        vocab[token_id] = bytes(value)
                    elif isinstance(value, str):
                        vocab[token_id] = value
                    else:
                        raise ValueError(f"Invalid vocab value type for token {key}: {type(value)}")

                raw_merged = state.get("merged_tokens", {})
                merged_tokens = {}
                for pair_key, new_token in raw_merged.items():
                    left, right = pair_key.split(",", 1)
                    merged_tokens[(int(left), int(right))] = int(new_token)

                self.__vocab = vocab if vocab else {0: "<|PAD|>"}
                normalization = state.get("normalization", {})
                self.normalization = {"lowercase": bool(normalization.get("lowercase", False))}
                self.merged_tokens = merged_tokens
                self.vocab_size = int(state.get("vocab_size", len(self.__vocab)))
                # Backward compat: old checkpoints don't have space_strategy field
                loaded_space_strategy = state.get("space_strategy", None)
                if loaded_space_strategy is None:
                    print("Warning: checkpoint missing 'space_strategy' — defaulting to 'gpt2' (legacy).")
                    loaded_space_strategy = SPACE_STRATEGY_GPT2
                self.space_strategy = loaded_space_strategy
                # Sync compiled_pattern with the resolved space strategy.
                self.compiled_pattern = (
                    _COMPILED_PATTERN_GPT2
                    if self.space_strategy == SPACE_STRATEGY_GPT2
                    else _COMPILED_PATTERN_METASPACE
                )
                self._reset_token_lookup_cache()
                return

            if extension == ".pkl":
                if not allow_pickle:
                    raise ValueError(
                        "Refusing to load pickle checkpoint without allow_pickle=True. "
                        "Use JSON checkpoint (kin_tokenizer.json) for safe loading."
                    )

                with open(resolved_path, "rb") as f:
                    tokenizer = pickle.load(f)
                    self.__vocab = tokenizer.__vocab
                    self.vocab_size = tokenizer.vocab_size
                    self.merged_tokens = tokenizer.merged_tokens
                    self.normalization = getattr(tokenizer, "normalization", {"lowercase": False})
                    self._reset_token_lookup_cache()
                return

            raise ValueError(f"Unsupported tokenizer checkpoint format: {resolved_path}")
        except FileNotFoundError:
            if not quiet_missing:
                logger.warning(f"Tokenizer file not found: {resolved_path}")
            return
        except (pickle.UnpicklingError, EOFError) as e:
            logger.error(f"Corrupted tokenizer file: {resolved_path} - {e}")
            raise
        except ValueError:
            raise
        except Exception as e:
            logger.error(f"Failed to load tokenizer: {e}")
            raise

    def get_tokens_pair_stats(self, tokens, stats=None):
        """
        Create frequencies of token pairs.
        tokens: list of tokens (int)
        stats: existing statistics dictionary to update
        """
        stats = {} if stats is None else stats
        # Inlined loop — avoids generator overhead for hot path
        n = len(tokens)
        for i in range(n - 1):
            pair = (tokens[i], tokens[i + 1])
            stats[pair] = stats.get(pair, 0) + 1
        return stats

    def create_tokens(self, text):
        """
        Create tokens from text.
        text: string of characters
        """
        if isinstance(text, str):
            text = text.encode("UTF-8")
        elif not isinstance(text, bytes):
            raise ValueError("Expected string or bytes")
        return list(map(int, text))

    def merge_tokens(self, pair, tokens, new_token):
        """
        Merge token pairs in a token list.
        pair: most frequent pair of tokens
        tokens: list of tokens
        new_token: the new token ID to replace the pair

        NOTE: This is a pure merge operation — it does NOT record the merge rule
        in self.merged_tokens. The caller (train()) is responsible for recording
        merge rules explicitly. FIX #10: Previously this had a side effect of
        writing self.merged_tokens[pair] = new_token, making encoding non-read-only.
        """
        new_tokens = []
        index = 0
        while index < len(tokens):
            if index < len(tokens) - 1 and pair[0] == tokens[index] and pair[1] == tokens[index + 1]:
                new_tokens.append(new_token)
                index += 2
            else:
                new_tokens.append(tokens[index])
                index += 1
        return new_tokens

    def remove_urls(self, text):
        """
        Remove URLs from a given text string.

        Args:
            text (str): The text to remove URLs from.

        Returns:
            str: The text with URLs removed.
        """
        url_pattern = re.compile(r'https?://\S+|www\.\S+')
        return url_pattern.sub(r'', text)

    # FIX #5: fix_case removed — vowel-capital regex no longer needed.

    def clean_lines(self, text):
        """
        Remove lines that contain only numbers, special characters,
        or mixes of numbers and special characters (no letters).
        """
        return '\n'.join(
            line for line in text.splitlines()
            if not regex.fullmatch(r'\s*([^\w\s]*\d+[^\w\s]*|[^\w\s]+|\d+)\s*', line)
        )

    def train(self, text, vocab_size=DEFAULT_VOCAB_SIZE, verbose=True, start_merge_iter=1, tokenizer_path=None, nbr_processes=None, lowercase=None):
        """
        Train the tokenizer using BPE.
        text: the text to be used for training the tokenizer
        vocab_size: the size of the vocabulary for the tokenizer after training
        start_merge_iter: The initial iteration/merge iteration
        tokenizer_path: The path where tokenizer will be saved (excluding filename)
        nbr_processes: Number of worker processes (default: cpu_count)
        """
        assert vocab_size >= MIN_VOCAB_SIZE
        if start_merge_iter > 1:
            append_logs = "a"
        else:
            append_logs = "w"
            self.reset_for_training()

        if tokenizer_path is not None:
            tokenizer_path = os.path.abspath(tokenizer_path)
            os.makedirs(tokenizer_path, exist_ok=True)

        # Keep retrain token-chunk state near the tokenizer checkpoint, not CWD-dependent.
        last_token_chunks_file_path = (
            _token_chunks_path(tokenizer_path)
            if tokenizer_path is not None
            else None
        )

        # check the vCPU(s)
        vcpus_nbr = _cpu_count()
        inline_text = f"AVAILABLE vCPU(s): {vcpus_nbr}\n\n"

        if nbr_processes is not None:
            vcpus_nbr = nbr_processes
        else:
            vcpus_nbr = _cpu_count()
            # For very large corpora, too many spawn workers on macOS can inflate memory.
            # Cap defaults conservatively unless user explicitly sets --num_processes.
            if len(text) >= (300 * 1024 * 1024):
                vcpus_nbr = min(vcpus_nbr, LARGE_CORPUS_CPU_CAP)

        if len(text) < 129:
            vcpus_nbr = 1
        elif len(text) <= (128 * vcpus_nbr):
            vcpus_nbr = len(text) // 128

        inline_text += f"Running {vcpus_nbr} processes\n"

        if start_merge_iter > 1:
            inline_text += f"Retraining started: {datetime.datetime.now(datetime.timezone.utc).strftime('%d-%m-%Y %H:%M:%S')} UTC\n\n"
        else:
            inline_text += f"Training started: {datetime.datetime.now(datetime.timezone.utc).strftime('%d-%m-%Y %H:%M:%S')} UTC\n\n"

        print(inline_text)

        start_time = time.time()
        print("\nInitial data cleaning using different regular expression patterns\n")
        inline_text += "\nInitial data cleaning using different regular expression patterns\n"
        lowercase = self._should_lowercase(lowercase)
        self.normalization["lowercase"] = lowercase

        # Unicode NFC normalization BEFORE any regex processing.
        # This ensures precomposed vs decomposed representations are unified
        # so BPE learns a single token for each visual character.
        text = unicodedata.normalize("NFC", text)

        if _HAS_RUST_MERGER and rust_preprocess_text is not None:
            # ── Rust-accelerated preprocessing ─────────────────────────
            # All 12+ regex substitutions run in a single Rust call.
            # For 800 MB text this is 3-5× faster than Python regex and
            # avoids creating 12 intermediate copies of the full string.
            text = rust_preprocess_text(text, lowercase)
        else:
            # ── Python fallback preprocessing ──────────────────────────
            text = regex.sub("\u2018", "'", text)
            text = regex.sub("\u2019", "'", text)
            text = regex.sub("\u201c", '"', text)
            text = regex.sub("\u201d", '"', text)
            # FIX #1: Removed circumflex stripping (â→a etc.) — these appear in
            # French loan words common in Kinyarwanda and must be preserved.
            text = self.clean_lines(text)
            text = regex.sub(r'(\n){3,}', '\n\n', text)
            # FIX #6: Strip trailing whitespace per line without doubling newlines.
            text = regex.sub(r'[^\S\n]+\n', '\n', text)
            # FIX #5: Removed vowel-capital regex r'([aeiou])([A-Z])' — English
            # camelCase heuristic that damages Kinyarwanda proper nouns.
            text = regex.sub(r'^(?!\s*$)\s+', '', text, flags=regex.MULTILINE)
            # FIX: Remove URLs BEFORE punctuation-joining to match Rust order.
            text = self.remove_urls(text)
            text = regex.sub(r'\*', ' ', text)
            text = regex.sub(r'[^\S\r\n]+', ' ', text)
            # FIX #7: Only glue common sentence/clause punctuation to prev word.
            text = regex.sub(r'(\w)\s+([.,;:!?])', r'\1\2', text)
            if lowercase:
                text = text.lower()

        total_time = time.time() - start_time
        print(f"\nInitial data cleaning using different regular expression patterns took: {total_time // 60} min {round(total_time % 60)} seconds\n")
        inline_text += f"\nInitial data cleaning using different regular expression patterns took: {total_time // 60} min {round(total_time % 60)} seconds\n"

        # FIX BUG-2: Use get_context() instead of set_start_method() to avoid RuntimeError on repeated calls.
        # set_start_method() can only be called once per process; get_context() is safe to call multiple times.
        mp_context = _get_mp_context()

        if start_merge_iter < 2:  # First training
            # Reading non kinyarwanda
            inline_text += "\nReading non kinyarwanda words\n"
            print("\nReading non kinyarwanda words\n")

            non_kinyarwanda_words_json_obj = importlib.resources.files('kin_tokenizer').joinpath('data/non_kinyarwanda_words.json')
            try:
                with importlib.resources.as_file(non_kinyarwanda_words_json_obj) as json_path:
                    with open(json_path, "r", encoding="utf-8") as f:
                        non_kinyarwanda_words = json.load(f)
            except FileNotFoundError:
                # Fall back to legacy pickle if JSON not found
                non_kinyarwanda_words_path_obj = importlib.resources.files('kin_tokenizer').joinpath('data/non_kinyarwanda_words.pkl')
                with importlib.resources.as_file(non_kinyarwanda_words_path_obj) as non_kinyarwanda_words_file_path:
                    with open(non_kinyarwanda_words_file_path, "rb") as f:
                        non_kinyarwanda_words = pickle.load(f)

            # Adding other initial values into a vocabulary
            self.ensure_base_byte_vocab()

            # FIX #3 + #17: Use compiled_pattern directly on the cleaned text,
            # matching exactly what encode() does.  This ensures BPE learns
            # merge rules on the same token atom shapes it will see at encode time.
            # METASPACE: prepend ▁ and replace spaces so every word — including
            # sentence-initial words — gets a consistent ▁ prefix before BPE sees it.
            # Only applies when space_strategy=="metaspace" (set in __init__; retraining
            # an old checkpoint inherits "gpt2" from load() and skips this step).
            if self.space_strategy == "metaspace":
                text = (" " + text.lstrip(" ")).replace(" ", METASPACE)
            print("\nPre-tokenizing text with compiled pattern (aligned with encode path)\n")
            inline_text += "\nPre-tokenizing text with compiled pattern (aligned with encode path)\n"
            start_time = time.time()
            if _HAS_RUST_MERGER and rust_pretokenize_text is not None:
                chunks = rust_pretokenize_text(text)
            else:
                chunks = regex.findall(self.compiled_pattern, text)
            # PERF: Free the original text immediately after splitting.
            del text
            gc.collect()
            total_time = time.time() - start_time
            print(f"\nPre-tokenizing text took: {total_time // 60} min {round(total_time % 60)} seconds\n")
            inline_text += f"\nPre-tokenizing text took: {total_time // 60} min {round(total_time % 60)} seconds\n"

            # removing non kinyarwanda words — filter checks the stripped chunk
            # against the word list so that space-prefixed chunks like " Aaron" still match "Aaron".
            print("\nRemoving non kinyarwanda words\n")
            inline_text += "\nRemoving non kinyarwanda words\n"
            start_time = time.time()
            non_kin_set = set(non_kinyarwanda_words)
            # FIX #18: Protect common French/English loan words used in Kinyarwanda.
            # The 244K-word filter list includes words like "police", "radio",
            # "president" that are actively used in Kinyarwanda text and should
            # NOT be filtered out. This allowlist overrides the filter.
            non_kin_set -= LOAN_WORD_ALLOWLIST
            # FIX #13: When training in lowercase mode, also lowercase the filter
            # set so that "Aaron" in the list matches "aaron" in the lowered text.
            if lowercase:
                lower_allowlist = {w.lower() for w in LOAN_WORD_ALLOWLIST}
                non_kin_set = {w.lower() for w in non_kin_set} - lower_allowlist
            # Strip ▁ prefix before comparing with the bare-word filter list.
            chunks = [c for c in chunks if c.lstrip(METASPACE).strip() not in non_kin_set]
            del non_kin_set
            total_time = time.time() - start_time

            print(f"\nRemoving non kinyarwanda words took: {total_time // 60} min {round(total_time % 60)} seconds\n")
            inline_text += f"\nRemoving non kinyarwanda words took: {total_time // 60} min {round(total_time % 60)} seconds\n"

            # removing existing words in the vocabulary
            print("\nRemoving words in chunks already exists in the tokenizer vocabulary\n")
            inline_text += "\nRemoving words in chunks already exists in the tokenizer vocabulary\n"
            start_time = time.time()

            existing_words_in_vocab = {value.decode("UTF-8", errors="ignore") for _, value in self.__vocab.items() if isinstance(value, bytes)}
            chunks = [c for c in chunks if c not in existing_words_in_vocab]
            del existing_words_in_vocab

            total_time = time.time() - start_time

            print(f"\nRemoving words in chunks already exists in the tokenizer vocabulary took: {total_time // 60} min {round(total_time % 60)} seconds\n")
            inline_text += f"\nRemoving words in chunks already exists in the tokenizer vocabulary took: {total_time // 60} min {round(total_time % 60)} seconds\n"

        else:  # Retraining
            print(f"\nLoading tokens chunks from a file: {last_token_chunks_file_path}\n")
            legacy_pkl_path = os.path.join(tokenizer_path, "current_token_chunks" + _TOKEN_CHUNKS_LEGACY_EXT)
            if not os.path.isfile(last_token_chunks_file_path):
                if os.path.isfile(legacy_pkl_path):
                    raise ValueError(
                        f"Only legacy pickle chunk file found: {legacy_pkl_path}. "
                        "Pickle token chunks are no longer loaded due to security risks. "
                        "Re-run first-train (without retrain) to generate a safe .bin checkpoint."
                    )
                raise FileNotFoundError(f"Token chunk checkpoint not found: {last_token_chunks_file_path}")
            tokens_chunks = _load_token_chunks(last_token_chunks_file_path)

            print("\nStarted retraining\n")
            inline_text += f"\nLoading tokens chunks from a file: {last_token_chunks_file_path}\n"
            inline_text += "\nStarted retraining\n"

        if start_merge_iter < 2:  # First training
            # FIX BUG-4: Use standalone function instead of creating a new KinTokenizer() per worker
            inline_text += "Creating tokens from text chunks..............\n"
            print("Creating tokens from text chunks..............\n")
            start_time = time.time()
            if _HAS_RUST_MERGER and rust_create_tokens_parallel is not None:
                # Rust+Rayon parallel UTF-8→u32 conversion — zero IPC/pickling
                tokens_chunks = rust_create_tokens_parallel(chunks)
            elif vcpus_nbr > 1:
                with mp_context.Pool(processes=vcpus_nbr) as pool:
                    tokens_chunks = pool.map(_create_tokens_standalone, chunks)
            else:
                tokens_chunks = list(map(_create_tokens_standalone, chunks))
            # Free the text chunks list immediately — tokens_chunks holds the
            # integer representation now. For an 800 MB corpus the string
            # chunks alone can occupy 1–2 GB.
            del chunks
            gc.collect()
            total_time = time.time() - start_time
            inline_text += f"Creating tokens took: {total_time // 60} min {round(total_time % 60)} seconds\n"
            print(f"Creating tokens took: {total_time // 60} min {round(total_time % 60)} seconds\n")

        num_merges = vocab_size - len(self.__vocab)  # We have encode tokens into range of 0 and 256

        # FIX BUG-11: Track next token ID as counter instead of O(V) max() per iteration
        next_token_id = max(self.__vocab.keys()) + 1

        # Keep logs in the same directory as the checkpoint for long-running jobs.
        log_file_path = os.path.join(tokenizer_path, "training_logs.txt") if tokenizer_path is not None else os.devnull

        # Clear pre-processing log accumulator — from here on we write directly
        # to the log file. Holding the entire log as a string wastes memory for
        # long-running 50K+ merge training sessions.
        del inline_text
        gc.collect()

        # Prefer Rust backend for merge loop speed when extension is available.
        # Falls back to Python multiprocessing loop if Rust extension is missing.
        if _HAS_RUST_MERGER and len(tokens_chunks) > 1 and num_merges > 0:
            # ── Batched Rust BPE training ──────────────────────────────────
            # Instead of sending ALL merges to Rust in a single blocking call
            # (which prevents progress output, checkpointing, and holds a full
            # copy of the corpus in Rust memory for the entire duration), we
            # call rust_bpe_train_loop in manageable batches.  Between batches
            # we: (a) print live progress, (b) save checkpoints, (c) let
            # Python's GC reclaim the previous chunk copies.
            # FIX #8: Increased batch size from 500 to 2000 to reduce the number
            # of Python↔Rust boundary copies. Each call serializes the full corpus,
            # so fewer calls = less total data transfer.
            RUST_BATCH_SIZE = RUST_MERGE_BATCH_SIZE
            total_tokens_approx = sum(len(c) for c in tokens_chunks)
            train_min_frequency = DEFAULT_MIN_FREQUENCY
            print(f"BPE min_frequency: {train_min_frequency} | Total tokens: {total_tokens_approx:,}")
            total_merges_done = 0
            merges_remaining = num_merges
            rust_overall_start = time.time()

            with open(log_file_path, append_logs) as f:
                while merges_remaining > 0:
                    batch_size = min(RUST_BATCH_SIZE, merges_remaining)
                    batch_start_time = time.time()

                    # ── Rust batch call ────────────────────────────────────
                    trained_chunks, merge_rules, _ = rust_bpe_train_loop(
                        tokens_chunks,
                        next_token_id,
                        batch_size,
                        train_min_frequency,
                    )

                    # Immediately replace the old Python list with the
                    # updated one and release the stale reference so GC can
                    # reclaim the previous ~N-hundred-MB allocation.
                    tokens_chunks = trained_chunks
                    del trained_chunks

                    batch_elapsed = time.time() - batch_start_time
                    actual_batch = len(merge_rules)

                    if actual_batch == 0:
                        # No more frequent pairs – training converged early
                        print(f"\nRust merge loop: no more pairs above min_frequency. Stopping early.")
                        f.write(f"\nRust merge loop: no more pairs above min_frequency. Stopping early.\n")
                        break

                    # ── Process merge rules from this batch ────────────────
                    for merge_offset, merge_rule in enumerate(merge_rules):
                        global_idx = start_merge_iter + total_merges_done + merge_offset
                        top_pair, new_token = merge_rule
                        self.merged_tokens[top_pair] = new_token
                        self.__vocab[new_token] = self.__vocab[top_pair[0]] + self.__vocab[top_pair[1]]

                        if verbose:
                            try:
                                decoded_pair = self.decode(list(top_pair), return_eos=False, multiprocess=False)
                            except Exception:
                                decoded_pair = "Error decoding"
                            line = (
                                f"Merge({global_idx}/{num_merges + start_merge_iter - 1}): Pair {top_pair} "
                                f"('{decoded_pair}') -> {new_token}. Vocab size: {len(self.__vocab)} "
                                f"completed at {datetime.datetime.now(datetime.timezone.utc).strftime('%d-%m-%Y %H:%M:%S')} UTC."
                            )
                            line_2 = "=" * len(line)
                            print(line)
                            print(line_2)
                            f.write(line + "\n" + line_2 + "\n")

                    # Advance counters
                    total_merges_done += actual_batch
                    self._reset_token_lookup_cache()
                    merges_remaining -= actual_batch
                    next_token_id += actual_batch

                    # ── Batch summary ──────────────────────────────────────
                    overall_elapsed = time.time() - rust_overall_start
                    merges_per_sec = total_merges_done / max(overall_elapsed, 0.001)
                    eta_seconds = merges_remaining / max(merges_per_sec, 0.001)
                    batch_summary = (
                        f"\n[Batch complete] {total_merges_done}/{num_merges} merges done "
                        f"({actual_batch} this batch in {batch_elapsed:.1f}s). "
                        f"Rate: {merges_per_sec:.1f} merges/s. "
                        f"ETA: {eta_seconds / 60:.1f} min remaining.\n"
                    )
                    print(batch_summary)
                    f.write(batch_summary + "\n")
                    f.flush()

                    # ── Checkpoint every 500 merges (aligned with batches) ─
                    if tokenizer_path is not None:
                        self.vocab_size = len(self.__vocab)
                        self.save(tokenizer_path)
                        ckpt_msg = f"Saving tokenizer at merge iteration: {start_merge_iter + total_merges_done}"
                        print(f"\n{ckpt_msg}\n")
                        f.write(f"\n{ckpt_msg}\n\n")

                        _save_token_chunks(tokens_chunks, last_token_chunks_file_path)
                        print(f"Saving tokens_chunks at merge iteration: {start_merge_iter + total_merges_done}\n")
                        f.write(f"Saving tokens_chunks at merge iteration: {start_merge_iter + total_merges_done}\n\n")

                    # ── Memory management between batches ──────────────────
                    # Force a GC cycle to reclaim stale chunk copies and any
                    # intermediate PyO3 conversion buffers.
                    gc.collect()

                rust_total_time = time.time() - rust_overall_start
                summary_line = (
                    f"Rust merge loop completed {total_merges_done} merges in "
                    f"{rust_total_time // 60:.0f} min {round(rust_total_time % 60)} seconds."
                )
                print(summary_line)
                f.write(summary_line + "\n")
        else:
            # Use the same min_frequency as the Rust path for consistency.
            train_min_frequency = DEFAULT_MIN_FREQUENCY
            # FIX BUG-3 (CRITICAL): Create Pool ONCE before the merge loop.
            # Previously, a new Pool was created on EVERY merge iteration (~50,000 times),
            # each time spawning/forking N worker processes. This was the single biggest
            # performance bottleneck, potentially adding days to training time.
            # FIX: Resolve rust_count_pairs ONCE before the loop instead of
            # importing inside the hot loop on every iteration.
            _rust_count_pairs = None
            if _HAS_RUST_MERGER and rust_bpe_train_loop is None:
                try:
                    from kin_merge import rust_count_pairs as _rust_count_pairs
                except ImportError:
                    _rust_count_pairs = None
            merge_pool = mp_context.Pool(processes=vcpus_nbr) if vcpus_nbr > 1 else None
            try:
                with open(log_file_path, append_logs) as f:
                    for idx in range(start_merge_iter, start_merge_iter + num_merges):
                        merge_start_time = time.time()
                        if len(tokens_chunks) > 1:
                            new_token = next_token_id

                            print("Calculating tokens pair statistics......")
                            # PERF: Use Rust pair counting when available even in Python merge loop
                            start_time = time.time()
                            if _rust_count_pairs is not None:
                                stats = _rust_count_pairs(tokens_chunks)
                            else:
                                stats = {}
                                for tokens in tokens_chunks:
                                    self.get_tokens_pair_stats(tokens, stats)

                            total_time = time.time() - start_time
                            print(f"Calculating tokens pair statistics took: {total_time // 60} min {round(total_time % 60)} seconds\n")

                            # Find the most frequent pair
                            top_pair = _select_top_pair(stats)

                            if stats.get(top_pair, 0) < train_min_frequency:  # no more frequent pairs
                                print(f"Stopping: top pair frequency {stats.get(top_pair, 0)} < min_frequency {train_min_frequency}")
                                break

                            # Parallel merging of tokens using the persistent pool
                            print("Creating merging arguments for parallel merging....")
                            start_time = time.time()
                            merge_args = [(top_pair, tokens, new_token) for tokens in tokens_chunks]
                            total_time = time.time() - start_time
                            print(f"Creating merging arguments for parallel merging took: {total_time // 60} min {round(total_time % 60)} seconds\n")

                            start_time = time.time()
                            print("Merging top pair tokens .....")
                            # FIX BUG-3 & BUG-4: Reuse persistent pool + standalone merge function
                            if merge_pool is not None:
                                tokens_chunks = merge_pool.map(_merge_tokens_standalone, merge_args)
                            else:
                                tokens_chunks = list(map(_merge_tokens_standalone, merge_args))

                            total_time = time.time() - start_time
                            print(f"Merging the top pair tokens took: {total_time // 60} min {round(total_time % 60)} seconds\n")

                            # Save the merge
                            self.merged_tokens[top_pair] = new_token

                            # Add new vocabulary into the vocab
                            self.__vocab[new_token] = self.__vocab[top_pair[0]] + self.__vocab[top_pair[1]]
                            next_token_id += 1

                            merge_total_time = time.time() - merge_start_time

                            if tokenizer_path is not None and ((idx + 1) % PYTHON_FALLBACK_CHECKPOINT_EVERY == 0):
                                self.vocab_size = len(self.__vocab)
                                self.save(tokenizer_path)
                                print(f"\nSaving tokenizer at merge iteration: {idx + 1}\n")
                                f.write(f"\nSaving tokenizer at merge iteration: {idx + 1}\n\n")

                                _save_token_chunks(tokens_chunks, last_token_chunks_file_path)
                                print(f"\nSaving tokens_chunks at merge iteration: {idx + 1}\n")
                                f.write(f"\nSaving tokens_chunks at merge iteration: {idx + 1}\n\n")

                            # print messages on the console
                            if verbose:
                                try:
                                    decoded_pair = self.decode(list(top_pair), return_eos=False, multiprocess=False)
                                except Exception:
                                    decoded_pair = "Error decoding"
                                line = f"Merge({idx}/{num_merges + start_merge_iter - 1}): Pair {top_pair} ('{decoded_pair}') -> {new_token} ({stats[top_pair]} occurrences). Vocab size: {len(self.__vocab)}. Iteration time: {merge_total_time // 60}m {round(merge_total_time % 60)}s completed at {datetime.datetime.now(datetime.timezone.utc).strftime('%d-%m-%Y %H:%M:%S')} UTC."
                                line_2 = "=" * len(line)
                                print(line)
                                print(line_2)
                                f.write(line + "\n" + line_2 + "\n")

                        else:
                            break  # no more pairs
            finally:
                if merge_pool is not None:
                    merge_pool.close()
                    merge_pool.join()

        # Adding special token(end of sequence)
        print("\nTraining finished. Adding special tokens...")
        self.ensure_special_tokens()
        self.vocab_size = len(self.__vocab)
        self._reset_token_lookup_cache()
        print("Special tokens added. Final vocabulary size:", self.vocab_size)

    def _encode_chunk(self, indexed_word):
        """
        Encode a single word or character(s).
        params:
            indexed_word: tuple of (index, word)
        """
        index, word = indexed_word

        tokens = self.create_tokens(word)
        # FIX #10: Use pure merge function instead of self.merge_tokens() which
        # has a side effect of writing to self.merged_tokens — encoding must be read-only.
        # FIX #9: Scan for min-rank pair directly instead of building a stats dict.
        merged = self.merged_tokens
        while len(tokens) > 1:
            best_pair = None
            best_rank = float("inf")
            for i in range(len(tokens) - 1):
                pair = (tokens[i], tokens[i + 1])
                rank = merged.get(pair, float("inf"))
                if rank < best_rank:
                    best_rank = rank
                    best_pair = pair
            if best_pair is None or best_rank == float("inf"):
                break
            tokens = _merge_tokens_standalone((best_pair, tokens, merged[best_pair]))

        return index, tokens

    # Removed deprecated alias _encode_chunck in v4; use _encode_chunk directly.

    def encode(self, text, nbr_processes=None, lowercase=None, _skip_normalize=False):
        """
        Encode text to tokens using the trained BPE merges.
        text: text to be encoded
        nbr_processes: number of worker processes
        _skip_normalize: internal flag to skip normalization (when caller already normalized)
        """
        if not _skip_normalize:
            text = self._normalize_text_for_encoding(text, lowercase)
        # Rust fast paths use the hardcoded metaspace RE_TOKEN_PATTERN.
        # For gpt2-mode checkpoints (backward compat) we must use the Python
        # compiled_pattern instead so that words without ▁ prefix are matched.
        _use_rust_pattern = (
            _HAS_RUST_MERGER
            and getattr(self, "space_strategy", DEFAULT_SPACE_STRATEGY) == SPACE_STRATEGY_METASPACE
        )
        if _use_rust_pattern and rust_encode_text is not None:
            # rust_encode_text returns list[int] via PyO3 — no need to re-wrap.
            # Now uses Rayon par_iter internally for parallel BPE encoding.
            return rust_encode_text(text, self.merged_tokens)

        if _use_rust_pattern and rust_pretokenize_text is not None:
            text_chunks = rust_pretokenize_text(text)
        else:
            text_chunks = regex.findall(self.compiled_pattern, text)  # Splitting text into chunks

        if not text_chunks:
            return []

        # FIX BUG-2/BUG-9: Use get_context() instead of set_start_method()
        mp_context = _get_mp_context()

        vcpus_nbr = _cpu_count()

        if nbr_processes is not None:
            vcpus_nbr = nbr_processes

        if len(text) < 129:  # each process must have at least 128 characters
            vcpus_nbr = 1
        elif len(text) <= (128 * vcpus_nbr):
            vcpus_nbr = len(text) // 128

        # Add indices to the chunks to preserve order
        indexed_chunks = [(i, chunk) for i, chunk in enumerate(text_chunks)]

        # Prefer Rust-accelerated batch encoding when available.
        # This runs BPE encoding for ALL words in parallel via Rayon,
        # avoiding Python multiprocessing IPC/pickling overhead entirely.
        if _HAS_RUST_MERGER and rust_encode_batch is not None and len(indexed_chunks) > 0:
            # BUG FIX: rust_encode_batch expects string words, not pre-indexed tuples with bytes.
            # Also convert merged_tokens values properly — Rust expects merge_rank → new_token mapping.
            result = rust_encode_batch(indexed_chunks, self.merged_tokens)
            result.sort(key=lambda x: x[0])
        elif vcpus_nbr > 1:
            # FIX PERF-ENCODE: Use pool initializer so merged_tokens is pickled
            # ONCE per worker process (at pool creation) rather than once per
            # task. For a 50K-token vocabulary this cuts IPC data transfer by
            # O(len(text_chunks)) × O(vocab_size).
            with mp_context.Pool(
                processes=vcpus_nbr,
                initializer=_init_encode_worker,
                initargs=(self.merged_tokens,),
            ) as pool:
                result = pool.map(_encode_chunk_standalone, indexed_chunks)
            result.sort(key=lambda x: x[0])
        else:
            result = list(map(self._encode_chunk, indexed_chunks))

        return [token for _, chunk_tokens in result for token in chunk_tokens]

    def encode_batch(self, texts, max_length=None, padding=False, truncation=False,
                     return_attention_mask=False, return_tensors=None, lowercase=None):
        """
        Encode a batch of texts with optional padding and truncation.

        Args:
            texts: list of strings to encode
            max_length: maximum sequence length (required if truncation=True or padding='max_length')
            padding: False, True/'longest', or 'max_length'
            truncation: whether to truncate sequences exceeding max_length
            return_attention_mask: whether to return attention masks (1 for real tokens, 0 for padding)
            return_tensors: None for lists, 'np' for numpy arrays

        Returns:
            dict with 'input_ids' and optionally 'attention_mask'
        """
        if not isinstance(texts, (list, tuple)):
            raise ValueError("Expected a list of strings")

        normalized_texts = [self._normalize_text_for_encoding(t, lowercase) for t in texts]

        # Only use rust_encode_texts when space_strategy is metaspace, because
        # the Rust encode_text_ids path uses the hardcoded metaspace RE_TOKEN_PATTERN.
        # For gpt2-mode legacy checkpoints we must fall through to the Python path
        # which respects self.compiled_pattern.
        _use_rust_encode = (
            _HAS_RUST_MERGER
            and rust_encode_texts is not None
            and getattr(self, "space_strategy", DEFAULT_SPACE_STRATEGY) == SPACE_STRATEGY_METASPACE
        )
        if _use_rust_encode:
            encoded = [list(map(int, seq)) for seq in rust_encode_texts(normalized_texts, self.merged_tokens)]
        else:
            # FIX #15: Pass _skip_normalize=True since texts are already normalized above.
            # This avoids a redundant O(n) NFC pass per text in the fallback path.
            encoded = [self.encode(t, _skip_normalize=True) for t in normalized_texts]

        # Truncation
        if truncation and max_length is not None:
            encoded = [seq[:max_length] for seq in encoded]

        # Padding
        pad_id = 0  # PAD token is always ID 0
        if padding:
            if padding == 'max_length':
                if max_length is None:
                    raise ValueError("max_length is required when padding='max_length'")
                target_len = max_length
            else:  # True or 'longest'
                target_len = max(len(seq) for seq in encoded) if encoded else 0

            attention_masks = []
            for i, seq in enumerate(encoded):
                pad_count = max(0, target_len - len(seq))
                attention_masks.append([1] * len(seq) + [0] * pad_count)
                encoded[i] = seq + [pad_id] * pad_count
        else:
            attention_masks = [[1] * len(seq) for seq in encoded]

        result = {'input_ids': encoded}
        if return_attention_mask:
            result['attention_mask'] = attention_masks

        if return_tensors == 'np':
            import numpy as np
            result = {k: np.array(v) for k, v in result.items()}

        return result

    def single_process_decode(self, indices, return_eos=True):
        """
        Decode tokens back to text (single process).
        indices: list of tokens to be decoded
        """
        if not isinstance(indices, (list, tuple)):
            raise ValueError("Expected list of integers")
        tokens = []
        has_eos = False
        for idx in indices:
            if idx not in self.__vocab:
                raise KeyError(f"Token {idx} does not exist in the vocabularies")
            if self.__vocab[idx] == "<|EOS|>" and return_eos:
                has_eos = True
                continue
            tokens.append(self.__vocab[idx])

        return self._compose_decoded_text(tokens, append_eos=has_eos and return_eos)

    def _compose_decoded_text(self, token_values, append_eos=False):
        """
        Decode a mixed sequence of bytes and special-string tokens while preserving order.
        """
        output_parts = []
        bytes_buffer = bytearray()

        for value in token_values:
            if isinstance(value, bytes):
                bytes_buffer.extend(value)
                continue

            if bytes_buffer:
                output_parts.append(bytes(bytes_buffer).decode("UTF-8", errors="ignore"))
                bytes_buffer.clear()

            output_parts.append(value if isinstance(value, str) else str(value))

        if bytes_buffer:
            output_parts.append(bytes(bytes_buffer).decode("UTF-8", errors="ignore"))

        result = "".join(output_parts)
        # METASPACE decode: only when this tokenizer was trained with metaspace strategy.
        if getattr(self, "space_strategy", DEFAULT_SPACE_STRATEGY) == SPACE_STRATEGY_METASPACE:
            result = result.replace(METASPACE, " ").lstrip(" ")
        if append_eos:
            result += "<|EOS|>"
        return result

    def multiprocess_decode_chunk(self, args):
        chunk, return_eos = args
        tokens_part = []
        has_eos = False
        invalid_indices = []
        for idx in chunk:
            if idx not in self.__vocab:
                invalid_indices.append(idx)
                continue
            if return_eos and self.__vocab[idx] == "<|EOS|>":
                has_eos = True
                continue
            tokens_part.append(self.__vocab[idx])
        return tokens_part, has_eos, invalid_indices

    def decode(self, indices, multiprocess=True, return_eos=True):
        """
        Decode token indices back to text.
        indices: list or tuple containing tokens indices
        multiprocess (bool): enable/disable multiprocessing (cpus > 1)
        return_eos: return EOS token
        """

        if not isinstance(indices, (list, tuple)):
            raise ValueError("Expected list of integers")

        # FIX PERF-5: Raised threshold from 65 to 10000.
        # Multiprocessing overhead (Pool creation, IPC, pickling) vastly outweighs
        # any parallel benefit for small token lists. 65 tokens is far too low.
        if multiprocess and len(indices) > DECODE_MULTIPROCESS_THRESHOLD:
            num_processes = _cpu_count()

            if len(indices) <= (64 * num_processes):
                num_processes = max(1, len(indices) // 64)

            if num_processes == 1:
                return self.single_process_decode(indices, return_eos)

            chunk_size = max(1, len(indices) // num_processes)
            chunks = [indices[i:i + chunk_size] for i in range(0, len(indices), chunk_size)]

            mp_context = _get_mp_context()
            with mp_context.Pool(
                processes=num_processes,
                initializer=_init_decode_worker,
                initargs=(self._KinTokenizer__vocab,),
            ) as pool:
                results = pool.map(_decode_chunk_standalone, [(chunk, return_eos) for chunk in chunks])

            tokens = []
            has_eos = False
            invalid_indices = []

            for tokens_part, chunk_has_eos, chunk_invalid in results:
                tokens.extend(tokens_part)
                has_eos |= chunk_has_eos
                invalid_indices.extend(chunk_invalid)

            if invalid_indices:
                raise KeyError(f"Tokens {invalid_indices} do not exist in the vocabularies")

            return self._compose_decoded_text(tokens, append_eos=return_eos and has_eos)
        else:
            return self.single_process_decode(indices, return_eos)
