cipher.py

"""Encrypt and decrypt text using the trigram dual-context rank cipher."""

import numpy as np
from scipy.optimize import linear_sum_assignment

import common
import lang as _lang

# Tunable defaults; each language may override via the "cipher" block in its
# config.json (lowercase keys, e.g. "bigram_weight"). Values below are tuned
# for English; sv and no carry their own sweep-tuned overrides.
VOWEL_BOOST = 3  # boost factor for vowels after 2 consecutive consonants
STREAK_BOOST = 1.5  # boost factor for streak-break at word boundaries
BIGRAM_WEIGHT = 0  # exponent for the mid-word/start pair-frequency penalty (0=off).
# The Markov levels already encode pair statistics, so with the trie active this
# double-counting hurts more than it helps for corpus-rich languages; Latin (smallest
# corpus, weakest Markov levels) still wants 0.5
END_BIGRAM_WEIGHT = 2  # pair-frequency exponent at the final position; None follows
# BIGRAM_WEIGHT. Split from the mid-word weight because the end tables are the only
# end-pair awareness in the scoring and must survive when BIGRAM_WEIGHT drops to 0
PAIR_FLOOR = 0.005  # floor for normalized pair frequency so unseen pairs scale, not zero out
DOUBLE_DEMOTE = 0.5  # extra demotion for unsupported doubled letters (reads as a typo)
LENGTH_START_WEIGHT = 1  # exponent for length-conditioned start letter modifier (0=off)
LENGTH_END_WEIGHT = 0  # exponent for length-conditioned end letter modifier (0=off)
MAX_LEN_BUCKET = 7  # word lengths 7+ share one bucket
LETTER_DEMOTE: dict[str, float] = {}  # per-letter score multiplier, for letters with very
# narrow follower support (Latin q -> u only); emitting them rarely beats emitting them wrong
SELF_MAP_PENALTY = 1000  # cost added when a letter would map to itself; >2 forbids all
# self-maps, values in (1, 2) permit rank-aligned mid-word self-maps, <1 collapses the
# cipher to identity at equal contexts (word starts), i.e. plaintext passes through
TRIE_WEIGHT = 0.95  # blend weight for the lexicon-trie level (0=off): scores become
# (1-w)*markov + w*P(letter | own word prefix, exact word length), steering short
# words toward completing real wordlist words. Off-trie prefixes fall back to markov.
TRIE_MAX_LEN = 6  # lexicon-trie applies to words this long or shorter; longer words
# can't realistically land on real words and per-word cache keys would thrash the cache
CONT_WEIGHT = 0.5  # exponent for continuation-breadth demotion (0=off): demote candidates
# whose (prev, candidate) pair has follower mass concentrated in one letter — they paint
# the next position into a corner (generalizes Latin's hand-tuned q demote)
PENULT_WEIGHT = 0.5  # exponent for penultimate endability (0=off): at the second-to-last
# position, weight candidates by how much word-ending mass follows them in the corpus.
# With the mid-word pair penalty at 0 it repairs long-word endings for en too; no keeps
# it off (endings already near-perfect, not worth the MSE)

_TUNABLE_NAMES = (
    "VOWEL_BOOST",
    "STREAK_BOOST",
    "BIGRAM_WEIGHT",
    "END_BIGRAM_WEIGHT",
    "PAIR_FLOOR",
    "DOUBLE_DEMOTE",
    "LENGTH_START_WEIGHT",
    "LENGTH_END_WEIGHT",
    "LETTER_DEMOTE",
    "TRIE_WEIGHT",
    "TRIE_MAX_LEN",
    "CONT_WEIGHT",
    "PENULT_WEIGHT",
)

# Per-language caches for pair-frequency matrices and resolved tunables.
_pair_freq_cache: dict[str, dict[str, dict]] = {}
_tunables_cache: dict[str, dict] = {}


def _tunables():
    """Resolve tunable constants for the current language.

    Module-level defaults, overridden by the optional "cipher" object in the
    language's config.json (keys are the lowercase constant names).
    """
    code = _lang.get_lang()
    if code not in _tunables_cache:
        overrides = _lang.config().get("cipher", {})
        _tunables_cache[code] = {
            name: overrides.get(name.lower(), globals()[name]) for name in _TUNABLE_NAMES
        }
    return _tunables_cache[code]


def _get_pair_freq_set():
    """Get or build the per-language pair-frequency matrices (mid/start/start_tri/
    end/end_tri/cont/endable)."""
    code = _lang.get_lang()
    if code not in _pair_freq_cache:
        _pair_freq_cache[code] = {}
    return _pair_freq_cache[code]


def _normalize_pair_counts(pair_counts, LETTERS):
    """Normalize pair counts per row: max follower → 1.0, floored at PAIR_FLOOR.

    The floor keeps unseen pairs as a small multiplier instead of zero, so the
    relative order of the underlying smoothed Markov scores survives the
    penalty and impossible pairs don't collapse into frequency-ordered ties.
    """
    floor = _tunables()["PAIR_FLOOR"]
    return {
        key: {y: max(row[y] / (max(row.values()) or 1), floor) for y in LETTERS}
        for key, row in pair_counts.items()
    }


def _get_pair_freq(bigram_transitions):
    """Compute normalized mid-word pair-frequency matrix from bigram transitions."""
    cache = _get_pair_freq_set()
    if "mid" in cache:
        return cache["mid"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    pair_counts = {x: dict.fromkeys(LETTERS, 0) for x in LETTERS}
    for ctx, nexts in bigram_transitions.items():
        if ctx.startswith("^"):
            continue
        prev = ctx[-1]
        if prev not in ALPHABET:
            continue
        for ch, count in nexts.items():
            if ch in ALPHABET:
                pair_counts[prev][ch] += count
    cache["mid"] = _normalize_pair_counts(pair_counts, LETTERS)
    return cache["mid"]


def _get_start_pair_freq(bigram_transitions):
    """Word-initial bigram pair freq: ^X→Y (1-char key)."""
    cache = _get_pair_freq_set()
    if "start" in cache:
        return cache["start"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    pair_counts = {x: dict.fromkeys(LETTERS, 0) for x in LETTERS}
    for ctx, nexts in bigram_transitions.items():
        if not ctx.startswith("^") or len(ctx) < 2:
            continue
        first = ctx[-1]
        if first not in ALPHABET:
            continue
        for ch, count in nexts.items():
            if ch in ALPHABET:
                pair_counts[first][ch] += count
    cache["start"] = _normalize_pair_counts(pair_counts, LETTERS)
    return cache["start"]


def _get_start_trigram_pair_freq(trigram_transitions):
    """Word-initial trigram pair freq: ^XY→Z (2-char key)."""
    cache = _get_pair_freq_set()
    if "start_tri" in cache:
        return cache["start_tri"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    pair_counts = {}
    for ctx, nexts in trigram_transitions.items():
        if not ctx.startswith("^") or len(ctx) != 3:
            continue
        key = ctx[1:]
        if not all(c in ALPHABET for c in key):
            continue
        if key not in pair_counts:
            pair_counts[key] = dict.fromkeys(LETTERS, 0)
        for ch, count in nexts.items():
            if ch in ALPHABET:
                pair_counts[key][ch] += count
    cache["start_tri"] = _normalize_pair_counts(pair_counts, LETTERS)
    return cache["start_tri"]


def _get_end_pair_freq(bigram_transitions):
    """Word-final bigram pair freq: X→Y$ (1-char key)."""
    cache = _get_pair_freq_set()
    if "end" in cache:
        return cache["end"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    pair_counts = {x: dict.fromkeys(LETTERS, 0) for x in LETTERS}
    for ctx, nexts in bigram_transitions.items():
        if len(ctx) != 2 or "$" not in nexts:
            continue
        x, y = ctx[0], ctx[1]
        if x in ALPHABET and y in ALPHABET:
            pair_counts[x][y] += nexts["$"]
    cache["end"] = _normalize_pair_counts(pair_counts, LETTERS)
    return cache["end"]


def _get_end_trigram_pair_freq(trigram_transitions):
    """Word-final trigram pair freq: XY→Z$ (2-char key)."""
    cache = _get_pair_freq_set()
    if "end_tri" in cache:
        return cache["end_tri"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    pair_counts = {}
    for ctx, nexts in trigram_transitions.items():
        if "$" not in nexts or len(ctx) != 3:
            continue
        if not all(c in ALPHABET for c in ctx):
            continue
        key = ctx[:2]
        last = ctx[2]
        if key not in pair_counts:
            pair_counts[key] = dict.fromkeys(LETTERS, 0)
        pair_counts[key][last] += nexts["$"]
    cache["end_tri"] = _normalize_pair_counts(pair_counts, LETTERS)
    return cache["end_tri"]


_trie_cache: dict[str, dict] = {}


def _get_prefix_trie():
    """(prefix, word_length) -> letter mass: which letters continue the prefix
    toward a wordlist word of exactly that length.

    The ("", L) rows reduce to the length-conditioned start table; deeper rows
    are what let the cipher complete real words instead of merely plausible
    letter sequences.
    """
    code = _lang.get_lang()
    if code not in _trie_cache:
        ALPHABET = common.alphabet()
        table: dict[tuple[str, int], dict[str, float]] = {}
        for word, freq in common.load_wordlist():
            if len(word) < 2 or any(c not in ALPHABET for c in word):
                continue
            for i, ch in enumerate(word):
                row = table.setdefault((word[:i], len(word)), {})
                row[ch] = row.get(ch, 0) + freq
        _trie_cache[code] = table
    return _trie_cache[code]


def _blend_trie(counts, prefix, alen, LETTERS):
    """Blend the lexicon-trie level over markov scores: (1-w)*markov + w*trie.

    Both rankings get this with their own prefix, so the modifier stays
    symmetric. No trie row (off-lexicon prefix) leaves the scores untouched.
    """
    row = _get_prefix_trie().get((prefix, alen))
    if not row:
        return
    w = _tunables()["TRIE_WEIGHT"]
    n = sum(row.values())
    for ch in LETTERS:
        counts[ch] = (1 - w) * counts[ch] + w * row.get(ch, 0) / n


def _get_continuation_freq(bigram_transitions):
    """Continuation breadth per (prev, candidate): letter-follower mass of the
    pair beyond its single best follower, normalized per prev row.

    Near-zero means the candidate forces the next position into one specific
    letter (Latin q -> u), so any non-top assignment there reads as a typo.
    Breadth, not volume: "aq" is frequent but has one follower.
    """
    cache = _get_pair_freq_set()
    if "cont" in cache:
        return cache["cont"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    rows = {}
    for prev in [*LETTERS, "^"]:
        row = {}
        for c in LETTERS:
            nexts = bigram_transitions.get(prev + c, {})
            counts = [cnt for z, cnt in nexts.items() if z in ALPHABET]
            row[c] = sum(counts) - max(counts) if counts else 0
        rows[prev] = row
    cache["cont"] = _normalize_pair_counts(rows, LETTERS)
    return cache["cont"]


def _get_endable(bigram_transitions):
    """Endability per candidate: corpus mass of word endings where the
    candidate is the second-to-last letter (sum over y of count("cy" -> $)),
    normalized like a pair-freq row. High means the letter commonly sits one
    position before a word end."""
    cache = _get_pair_freq_set()
    if "endable" in cache:
        return cache["endable"]
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    mass = dict.fromkeys(LETTERS, 0)
    for ctx, nexts in bigram_transitions.items():
        if len(ctx) == 2 and "$" in nexts and ctx[0] in ALPHABET and ctx[1] in ALPHABET:
            mass[ctx[0]] += nexts["$"]
    cache["endable"] = _normalize_pair_counts({"pen": mass}, LETTERS)["pen"]
    return cache["endable"]


_length_tables_cache: dict[str, tuple[dict, dict]] = {}


def _get_length_tables():
    """Per-length-bucket start/end letter multipliers, built from the wordlist.

    First and last letters distribute differently by word length (2-letter
    words start/end very differently from 7-letter ones). Buckets are exact
    lengths 2..MAX_LEN_BUCKET-1 with MAX_LEN_BUCKET collecting the rest.
    """
    code = _lang.get_lang()
    if code not in _length_tables_cache:
        LETTERS = common.letters()
        start_counts: dict[int, dict[str, int]] = {}
        end_counts: dict[int, dict[str, int]] = {}
        for word, freq in common.load_wordlist():
            if len(word) < 2:
                continue
            bucket = min(len(word), MAX_LEN_BUCKET)
            srow = start_counts.setdefault(bucket, dict.fromkeys(LETTERS, 0))
            erow = end_counts.setdefault(bucket, dict.fromkeys(LETTERS, 0))
            if word[0] in srow:
                srow[word[0]] += freq
            if word[-1] in erow:
                erow[word[-1]] += freq
        _length_tables_cache[code] = (
            _normalize_pair_counts(start_counts, LETTERS),
            _normalize_pair_counts(end_counts, LETTERS),
        )
    return _length_tables_cache[code]


def hungarian_derangement(p_ranking, c_ranking):
    """Optimal rank-to-rank assignment with self-map penalty.

    For squared rank displacement (a convex cost), the optimal assignment
    between two rankings is the rank-k-to-rank-k pairing, locally perturbed
    only where the self-map penalty forces swaps. Cost and tiebreaker both
    transpose when the contexts swap, which the involution property requires.
    """
    n = len(p_ranking)
    cost = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            cost[i, j] = (i - j) ** 2 + 0.001 * i * j
            if p_ranking[i] == c_ranking[j]:
                cost[i, j] += SELF_MAP_PENALTY
    row_ind, col_ind = linear_sum_assignment(cost)
    perm = {p_ranking[i]: c_ranking[j] for i, j in zip(row_ind, col_ind, strict=True)}

    # For identical rankings with odd alphabet, the Hungarian result can't be
    # an involution (no fixed-point-free involution exists for odd n). Fix by
    # letting the rarest letter in the odd cycle map to itself.
    if p_ranking == c_ranking and n % 2 == 1:
        rank_of = {ch: i for i, ch in enumerate(p_ranking)}
        visited = set()
        for start in p_ranking:
            if start in visited:
                continue
            cycle = []
            ch = start
            while ch not in visited:
                visited.add(ch)
                cycle.append(ch)
                ch = perm[ch]
            if len(cycle) > 2:
                rarest_idx = max(range(len(cycle)), key=lambda k: rank_of[cycle[k]])
                rarest = cycle[rarest_idx]
                remaining = cycle[:rarest_idx] + cycle[rarest_idx + 1 :]
                perm[rarest] = rarest
                for k in range(0, len(remaining), 2):
                    perm[remaining[k]] = remaining[k + 1]
                    perm[remaining[k + 1]] = remaining[k]

    return perm


def _apply_context_modifiers(
    counts,
    ctx,
    bigram_transitions,
    trigram_transitions,
    word_end,
    length_bucket=None,
    word_penult=False,
):
    """Apply pair-frequency penalty, length-conditioned boundary weighting, and
    vowel boost to Markov counts based on context."""
    LETTERS = common.letters()
    ALPHABET = common.alphabet()
    VOWEL_SET = common.vowel_set()
    tun = _tunables()
    prev = ctx[-1]
    end_weight = tun["END_BIGRAM_WEIGHT"]
    if end_weight is None:
        end_weight = tun["BIGRAM_WEIGHT"]
    pair_weight = end_weight if word_end else tun["BIGRAM_WEIGHT"]
    if pair_weight and prev in ALPHABET:
        at_word_start = ctx.startswith("^")
        pair_freq = None
        key = None

        if at_word_start:
            if len(ctx) == 3:
                pf = _get_start_trigram_pair_freq(trigram_transitions)
                tri_key = ctx[1:]
                if tri_key in pf:
                    pair_freq, key = pf, tri_key
            if pair_freq is None:
                pair_freq = _get_start_pair_freq(bigram_transitions)
                key = prev
        elif word_end:
            prev2 = ctx[-2:]
            if len(prev2) == 2 and all(c in ALPHABET for c in prev2):
                pf = _get_end_trigram_pair_freq(trigram_transitions)
                if prev2 in pf:
                    pair_freq, key = pf, prev2
            if pair_freq is None:
                pair_freq = _get_end_pair_freq(bigram_transitions)
                key = prev
        else:
            pair_freq = _get_pair_freq(bigram_transitions)
            key = prev

        for ch in LETTERS:
            pf = pair_freq[key][ch]
            if ch == prev and pf <= tun["PAIR_FLOOR"]:
                pf *= tun["DOUBLE_DEMOTE"]
            counts[ch] *= pf**pair_weight

    if length_bucket is not None:
        start_tab, end_tab = _get_length_tables()
        if tun["LENGTH_START_WEIGHT"] and ctx == "^" and length_bucket in start_tab:
            row = start_tab[length_bucket]
            for ch in LETTERS:
                counts[ch] *= row[ch] ** tun["LENGTH_START_WEIGHT"]
        elif tun["LENGTH_END_WEIGHT"] and word_end and length_bucket in end_tab:
            row = end_tab[length_bucket]
            for ch in LETTERS:
                counts[ch] *= row[ch] ** tun["LENGTH_END_WEIGHT"]

    for ch, factor in tun["LETTER_DEMOTE"].items():
        if ch in counts:
            counts[ch] *= factor

    if tun["CONT_WEIGHT"] and not word_end:
        cont_row = _get_continuation_freq(bigram_transitions).get(prev)
        if cont_row:
            for ch in LETTERS:
                counts[ch] *= cont_row[ch] ** tun["CONT_WEIGHT"]

    if tun["PENULT_WEIGHT"] and word_penult:
        endable = _get_endable(bigram_transitions)
        for ch in LETTERS:
            counts[ch] *= endable[ch] ** tun["PENULT_WEIGHT"]

    last2 = ctx[-2:]
    after_cc = (
        len(last2) >= 2
        and last2[-1] in ALPHABET
        and last2[-1] not in VOWEL_SET
        and last2[-2] in ALPHABET
        and last2[-2] not in VOWEL_SET
    )
    if after_cc and tun["VOWEL_BOOST"] > 1:
        for ch in LETTERS:
            if ch in VOWEL_SET:
                counts[ch] *= tun["VOWEL_BOOST"]

    # Streak-break boost: prevent all-vowel or all-consonant short words
    if word_end and ctx.startswith("^"):
        since_start = [c for c in ctx[1:] if c in ALPHABET]
        if since_start:
            if all(c not in VOWEL_SET for c in since_start):
                for ch in LETTERS:
                    if ch in VOWEL_SET:
                        counts[ch] *= tun["STREAK_BOOST"]
            elif all(c in VOWEL_SET for c in since_start):
                for ch in LETTERS:
                    if ch not in VOWEL_SET:
                        counts[ch] *= tun["STREAK_BOOST"]


def _make_perm(
    bigram_transitions,
    trigram_transitions,
    plain_ctx,
    cipher_ctx,
    word_end=False,
    length_bucket=None,
    prefixes=None,
    word_penult=False,
):
    """Compute permutation for a (plain_ctx, cipher_ctx) pair using trigram context.

    prefixes, when set, is (plain_prefix, cipher_prefix, alpha_len) and enables
    the lexicon-trie level: each ranking is blended toward completing a real
    word from its own emitted prefix.
    """
    LETTERS = common.letters()
    feat = common.features()

    p_counts = common.markov_counts(LETTERS, bigram_transitions, plain_ctx, trigram_transitions)
    c_counts = common.markov_counts(LETTERS, bigram_transitions, cipher_ctx, trigram_transitions)

    if prefixes is not None:
        plain_prefix, cipher_prefix, alen = prefixes
        _blend_trie(p_counts, plain_prefix, alen, LETTERS)
        _blend_trie(c_counts, cipher_prefix, alen, LETTERS)

    _apply_context_modifiers(
        p_counts,
        plain_ctx,
        bigram_transitions,
        trigram_transitions,
        word_end,
        length_bucket,
        word_penult,
    )
    _apply_context_modifiers(
        c_counts,
        cipher_ctx,
        bigram_transitions,
        trigram_transitions,
        word_end,
        length_bucket,
        word_penult,
    )

    p_ranking = sorted(LETTERS, key=lambda ch: (-p_counts[ch], -feat[ch]["Frequency"], ch))
    c_ranking = sorted(LETTERS, key=lambda ch: (-c_counts[ch], -feat[ch]["Frequency"], ch))

    return hungarian_derangement(p_ranking, c_ranking)


def _get_perm(
    bigram_transitions,
    trigram_transitions,
    plain_ctx,
    cipher_ctx,
    cache,
    word_end=False,
    length_bucket=None,
    prefixes=None,
    word_penult=False,
):
    """Get permutation and its inverse, using runtime cache for repeated context pairs."""
    # Length conditioning only affects word-boundary positions; drop it from
    # the key (and the perm computation) elsewhere so mid-word perms stay shared
    # (when the trie is active, keys are per-prefix anyway and sharing is moot).
    if not (plain_ctx == "^" or word_end):
        length_bucket = None
    end_prefix = "end:" if word_end else ("pen:" if word_penult else "")
    # Trie-conditioned perms depend on the full prefixes and exact word length,
    # not just the 3-char contexts, so they get their own cache keys.
    trie_part = f"T{prefixes[2]}:{prefixes[0]}|{prefixes[1]}:" if prefixes is not None else ""
    cache_key = f"{end_prefix}{trie_part}L{length_bucket}:{plain_ctx}|{cipher_ctx}"
    if cache_key in cache:
        return cache[cache_key]
    perm = _make_perm(
        bigram_transitions,
        trigram_transitions,
        plain_ctx,
        cipher_ctx,
        word_end=word_end,
        length_bucket=length_bucket,
        prefixes=prefixes,
        word_penult=word_penult,
    )
    inv_perm = {v: k for k, v in perm.items()}
    cache[cache_key] = (perm, inv_perm)
    return cache[cache_key]


def encrypt_word(word, bigram_transitions, trigram_transitions, cache):
    """Encrypt a single word using trigram dual-context rank cipher, preserving case."""
    ALPHABET = common.alphabet()
    alen = common.alpha_len(word)

    # Single-letter words pass through unchanged: the set of real one-letter
    # words is tiny and closed (en {a, I}, sv {i, å, ö}), so any substitution
    # maps them out of it — and a fixed 1-letter mapping hides nothing anyway.
    if alen == 1:
        return word

    tun = _tunables()
    trie_on = tun["TRIE_WEIGHT"] > 0 and alen <= tun["TRIE_MAX_LEN"]

    plain_ctx = "^"
    cipher_ctx = "^"
    plain_prefix = ""
    cipher_prefix = ""
    result = []
    alpha_pos = 0

    for ch in word:
        low = ch.lower()
        if low not in ALPHABET:
            result.append(ch)
            continue

        alpha_pos += 1
        is_last = alpha_pos == alen
        is_penult = bool(tun["PENULT_WEIGHT"]) and alpha_pos == alen - 1
        prefixes = (plain_prefix, cipher_prefix, alen) if trie_on else None
        perm, _ = _get_perm(
            bigram_transitions,
            trigram_transitions,
            plain_ctx,
            cipher_ctx,
            cache,
            word_end=is_last,
            length_bucket=min(alen, MAX_LEN_BUCKET),
            prefixes=prefixes,
            word_penult=is_penult,
        )
        mapped = perm[low]
        result.append(mapped.upper() if ch.isupper() else mapped)

        plain_ctx = (plain_ctx + low)[-3:]
        cipher_ctx = (cipher_ctx + mapped)[-3:]
        plain_prefix += low
        cipher_prefix += mapped

    return "".join(result)


def decrypt_word(word, bigram_transitions, trigram_transitions, cache):
    """Decrypt a single word using trigram dual-context rank cipher, preserving case."""
    ALPHABET = common.alphabet()
    alen = common.alpha_len(word)

    # Single-letter words pass through unchanged (see encrypt_word).
    if alen == 1:
        return word

    tun = _tunables()
    trie_on = tun["TRIE_WEIGHT"] > 0 and alen <= tun["TRIE_MAX_LEN"]

    cipher_ctx = "^"
    plain_ctx = "^"
    plain_prefix = ""
    cipher_prefix = ""
    result = []
    alpha_pos = 0

    for ch in word:
        low = ch.lower()
        if low not in ALPHABET:
            result.append(ch)
            continue

        alpha_pos += 1
        is_last = alpha_pos == alen
        is_penult = bool(tun["PENULT_WEIGHT"]) and alpha_pos == alen - 1
        prefixes = (plain_prefix, cipher_prefix, alen) if trie_on else None
        _, inv_perm = _get_perm(
            bigram_transitions,
            trigram_transitions,
            plain_ctx,
            cipher_ctx,
            cache,
            word_end=is_last,
            length_bucket=min(alen, MAX_LEN_BUCKET),
            prefixes=prefixes,
            word_penult=is_penult,
        )
        mapped = inv_perm[low]
        result.append(mapped.upper() if ch.isupper() else mapped)

        cipher_ctx = (cipher_ctx + low)[-3:]
        plain_ctx = (plain_ctx + mapped)[-3:]
        cipher_prefix += low
        plain_prefix += mapped

    return "".join(result)


def _cipher_text(text, bigram_transitions, trigram_transitions, word_fn):
    """Process text word-by-word, splitting on non-alpha and preserving structure."""
    ALPHABET = common.alphabet()
    NON_ALPHA_RE = common.non_alpha_re()
    cache = {}
    tokens = NON_ALPHA_RE.split(text)
    return "".join(
        word_fn(tok, bigram_transitions, trigram_transitions, cache)
        if tok and any(c.lower() in ALPHABET for c in tok)
        else tok
        for tok in tokens
    )


def encrypt_text(text, bigram_transitions, trigram_transitions):
    """Encrypt text, splitting on non-alpha and preserving structure."""
    return _cipher_text(text, bigram_transitions, trigram_transitions, encrypt_word)


def decrypt_text(text, bigram_transitions, trigram_transitions):
    """Decrypt text, splitting on non-alpha and preserving structure."""
    return _cipher_text(text, bigram_transitions, trigram_transitions, decrypt_word)


if __name__ == "__main__":
    import argparse
    import sys

    parser = argparse.ArgumentParser(description="Trigram dual-context rank cipher")
    parser.add_argument("text", nargs="*", help="text to encrypt/decrypt")
    parser.add_argument("-d", "--decrypt", action="store_true", help="decrypt instead of encrypt")
    parser.add_argument("-i", "--input", help="input file (UTF-8)")
    parser.add_argument("-o", "--output", help="output file (UTF-8)")
    parser.add_argument(
        "--lang", required=True, choices=_lang.available_languages(), help="language code"
    )
    args = parser.parse_args()

    _lang.set_lang(args.lang)

    line_ending: bytes | None = None
    if args.input:
        with open(args.input, "rb") as f:
            raw = f.read()
        line_ending = b"\r\n" if b"\r\n" in raw else b"\n"
        text = raw.decode("utf-8")
        if text.endswith("\n"):
            text = text[:-1]
        if text.endswith("\r"):
            text = text[:-1]
        text = text.replace("\r\n", "\n")
    elif args.text:
        text = " ".join(args.text)
    elif not sys.stdin.isatty():
        text = sys.stdin.read().rstrip("\n")
    else:
        text = input("Enter text: ")

    bigram_transitions = common.load_bigram_transitions()
    trigram_transitions = common.load_trigram_transitions()

    if args.decrypt:
        result = decrypt_text(text, bigram_transitions, trigram_transitions)
    else:
        result = encrypt_text(text, bigram_transitions, trigram_transitions)

    if args.output:
        result += "\n"
        if line_ending == b"\r\n":
            result = result.replace("\n", "\r\n")
        with open(args.output, "wb") as out:
            out.write(result.encode("utf-8"))
    else:
        print(result)