"""
tokenmax.py — Deterministic token-maxing post-processor for compressed think blocks.

Applies ONLY substitutions that are verified to save tokens on the Qwen 248K tokenizer
(OmniCoder-9B / Qwen3.5). Every substitution was tested in-context (not isolation) to
confirm real token savings without boundary interference.

Design:
  - LLM does semantic compression (what to keep vs drop)
  - This code enforces consistent notation deterministically
  - GUARD: only returns the processed version if ntok(result) < ntok(original)
  - Idempotent: safe to run multiple times

Usage:
    from caveman.compress.tokenmax import tokenmax, tokenmax_with_stats
    compressed = tokenmax(think_text, tokenizer)
    compressed, stats = tokenmax_with_stats(think_text, tokenizer)

Verified: 2026-05-31 on Qwen 248K vocab. 27/28 substitutions save in-context.
Zero false positives. One zero-effect (贪心 for "greedy" — boundary-dependent).
"""

import re
from typing import Optional

# ── Phase 1: Filler drops ──────────────────────────────────────────────
# Phrases that carry zero information in compressed reasoning.
# Only patterns that are NEVER load-bearing in a think block.
_FILLER_PATTERNS = [
    # Metacognition (the model narrating its own process)
    r"\bI need to\b",
    r"\bwe need to\b",
    r"\bI will\b",
    r"\bI'll\b",
    r"\blet me\b",
    r"\blet's\b",
    r"\bI want to\b",
    r"\bI should\b",
    r"\bwe should\b",
    # Hedging
    r"\bprobably\b",
    r"\bbasically\b",
    r"\bessentially\b",
    r"\bit seems like\b",
    # Filler transitions
    r"\bin order to\b",
    r"\bfirst of all\b",
    r"\bin other words\b",
    r"\bon the other hand\b",
    r"\bmore specifically\b",
    r"\bto be more precise\b",
    r"\band so on\b",
    # Conversational padding (require word boundary at end to avoid "Greatest", "Perfectly")
    r"\bGreat\b[,!.]?\s*",
    r"\bPerfect\b[,!.]?\s*",
    # Obvious statements
    r"\bAs (?:we|you) can see\b",
    r"\bAs mentioned (?:above|earlier|before)\b",
]

# ── Phase 2: Phrase → cheapest token ───────────────────────────────────
# Ordered LONGEST FIRST to prevent partial matches.
# Each entry: (regex_pattern, replacement, category)
# Categories: 'cjk', 'symbol', 'abbrev' — for stats tracking.
_SUBSTITUTIONS = [
    # ── COMPOUND PATTERNS FIRST (must fire before their components) ──

    # Verbose comparison phrases (+5t savings)
    (r'\bis\s+greater\s+than\s+or\s+equal\s+to\b', '≥', 'symbol'),   # +5t
    (r'\bis\s+less\s+than\s+or\s+equal\s+to\b', '≤', 'symbol'),      # +5t

    # Verbose discourse phrases (+3t savings)
    (r'\bwe\s+can\s+see\s+that\b', '可知', 'cjk'),                    # +3t
    (r'\bat\s+the\s+same\s+time\b', '同时', 'cjk'),                   # +3t
    (r'\bthat\s+is\s+to\s+say\b', '即', 'cjk'),                      # +3t

    # Multi-word phrases (+2t savings)
    (r'\bin\s+this\s+case\b', '此时', 'cjk'),                         # +2t
    (r'\bthe\s+number\s+of\b', '个数', 'cjk'),                        # +2t
    (r'\bis\s+equal\s+to\b', '等于', 'cjk'),                          # +2t

    # Complexity boilerplate (biggest per-occurrence savings)
    (r'[Oo]\(n\)\s*time[,;]?\s*[Oo]\(n\)\s*space\.?', 'O(n|n).', 'abbrev'),
    (r'[Oo]\(n\)\s*time[,;]?\s*[Oo]\(1\)\s*space\.?', 'O(n|1).', 'abbrev'),
    (r'[Oo]\(n\s*log\s*n\)\s*time[,;]?\s*[Oo]\(n\)\s*space', 'O(n㏒n|n)', 'abbrev'),
    (r'[Oo]\(n\s*log\s*n\)\s*time[,;]?\s*[Oo]\(1\)\s*space', 'O(n㏒n|1)', 'abbrev'),
    (r'[Tt]ime\s*complexity[:\s]+', 'T=', 'abbrev'),
    (r'[Ss]pace\s*complexity[:\s]+', 'S=', 'abbrev'),

    # Multi-word compounds (BEFORE their single-word components)
    (r'\bassume without loss of generality\b', '设 不妨', 'cjk'),  # before "assume"
    (r'\bproof by contradiction\b', '反证', 'cjk'),                # before "proof", "contradiction"
    (r'\bnecessary and sufficient\b', '充要', 'cjk'),              # before "sufficient"
    (r'\bnot equal(?:\s+to)?\b', '≠', 'symbol'),                  # before "is not", "does not"
    (r'\bif and only if\b', 'iff', 'abbrev'),                     # before "for all"
    (r'\bmuch greater than\b', '≫', 'symbol'),                    # before "greater than"
    (r'\bkeep track(?:\s+of)?\b', '记录', 'cjk'),                 # before article strip
    (r'\bin ascending order\b', 'asc', 'abbrev'),
    (r'\bin descending order\b', 'desc', 'abbrev'),
    (r'\bmaximum value\b', '最大值', 'cjk'),
    (r'\bminimum value\b', '最小值', 'cjk'),
    (r'\breturn value\b', '返回値', 'cjk'),
    (r'\brather than\b', '而非', 'cjk'),
    (r'\baccording to\b', '按照', 'cjk'),

    # DS compounds (before components)
    (r'\bdoubly linked list\b', 'DLL', 'abbrev'),      # BEFORE "linked list"
    (r'\bbinary indexed tree\b', 'BIT', 'abbrev'),     # BEFORE "binary"
    (r'\bminimum spanning tree\b', 'MST', 'abbrev'),
    (r'\bdepth[- ]first search\b', 'DFS', 'abbrev'),
    (r'\bbreadth[- ]first search\b', 'BFS', 'abbrev'),
    (r'\bdynamic programming\b', 'DP', 'abbrev'),
    (r'\bdivide and conquer\b', '分治', 'cjk'),
    (r'\bmonot(?:onic|one)\s*stack\b', '单调栈', 'cjk'),
    (r'\btime limit exceeded\b', '超时', 'cjk'),
    (r'\bout of bounds\b', '越界', 'cjk'),
    (r'\bremove duplicates?\b', '去重', 'cjk'),
    (r'\benumerate all\b', '穷举', 'cjk'),
    (r'\bbinary search\b', '二分', 'cjk'),
    (r'\bsliding window\b', 'sw', 'abbrev'),
    (r'\bunion[- ]find\b', 'UF', 'abbrev'),
    (r'\btopological sort\b', '拓扑序', 'cjk'),
    (r'\bshortest path\b', 'sp', 'abbrev'),
    (r'\blinked list\b', 'LL', 'abbrev'),
    (r'\bpriority queue\b', 'heap', 'abbrev'),
    (r'\bprefix sum\b', 'ps', 'abbrev'),
    (r'\bbrute force\b', '暴力', 'cjk'),
    (r'\bno solution\b', '无解', 'cjk'),
    (r'\bedge cases?\b', '边界', 'cjk'),
    (r'\bbase case\b', 'bc', 'abbrev'),
    (r'\bworst case\b', 'wc', 'abbrev'),

    # ── SINGLE-WORD SUBSTITUTIONS (safe after compounds consumed) ──

    # +4t savings
    (r'\bobviously\b', '显然', 'cjk'),
    # +3t savings
    (r'\bredundant\b', '冗余', 'cjk'),
    (r'\bsatisf(?:y|ies|ied)\b', '满足', 'cjk'),            # +2t, 333x in data
    (r'\bunsorted\b', '无序', 'cjk'),
    (r'\bdue to\b', '由于', 'cjk'),
    (r'\bhence\b', '故', 'cjk'),
    (r'\bnamely\b', '即', 'cjk'),
    (r'\bassume\b', '设', 'cjk'),
    (r'\bsuppose\b', '设', 'cjk'),
    (r'\bderive\b', '推导', 'cjk'),
    # +2t savings
    (r'\bmonotone\b', '单调', 'cjk'),
    (r'\bconvergent\b', '收敛', 'cjk'),
    (r'\bdivergent\b', '发散', 'cjk'),
    (r'\bcommutative\b', '交换', 'cjk'),
    (r'\bdeterministic\b', '确定', 'cjk'),
    (r'\bprobabilistic\b', '概率', 'cjk'),
    (r'\bprove\b', '证明', 'cjk'),
    (r'\bproof\b', '证明', 'cjk'),
    (r'\bflip\b', '翻转', 'cjk'),
    (r'\bsorted\b(?!\s*[=(\[])', '有序', 'cjk'),  # not before = ( [ (assignment/call)
    # +1t savings
    (r'\bbacktrack(?:ing)?\b', '回溯', 'cjk'),
    (r'\btravers(?:e|al|ing)\b', '遍历', 'cjk'),
    (r'\brecursi(?:on|ve|vely)\b', '递归', 'cjk'),
    (r'\bcontradiction\b', '矛盾', 'cjk'),
    (r'\bsufficient\b', '充分', 'cjk'),
    (r'\bequivalent\b', '等价', 'cjk'),
    (r'\bsymmetric\b', '对称', 'cjk'),
    (r'\binvariant\b', '不变', 'cjk'),
    (r'\bexponential\b', '指数', 'cjk'),
    (r'\bpermutation\b', '排列', 'cjk'),
    (r'\badjacent\b', '相邻', 'cjk'),
    (r'\boptimal\b', '最优', 'cjk'),
    (r'\bfeasible\b', '可行', 'cjk'),
    (r'\binduction\b', '归纳', 'cjk'),
    (r'\bmaintain\b', '维护', 'cjk'),
    (r'\bswap\b(?!\s*[=(\[])', '交换', 'cjk'),  # not before = ( [ (assignment/call)
    (r'\bcumulative\b', '累积', 'cjk'),
    (r'\bquotient\b', '商', 'cjk'),
    (r'\bmemoiz(?:ation|e)\b', 'memo', 'abbrev'),
    # +2t savings (mined from v19 data)
    (r'\bmathematical\b', '数学', 'cjk'),
    (r'\bcorresponding(?:ly)?\b', '对应', 'cjk'),
    (r'\brequirement\b', '需求', 'cjk'),
    # +1t savings (mined from v19 data)
    (r'\bcomplexity\b', '复杂度', 'cjk'),
    (r'\bsimilarly\b', '同理', 'cjk'),
    (r'\bsubstitut(?:e|ion)\b', '代入', 'cjk'),
    (r'\bincreasing(?:ly)?\b', '递增', 'cjk'),
    (r'\bdecreasing(?:ly)?\b', '递减', 'cjk'),
    (r'\brespectively\b', '分别', 'cjk'),
    (r'\bnecessarily\b', '必然', 'cjk'),
    (r'\btransformation\b', '变换', 'cjk'),
    (r'\bprerequisite\b', '前提', 'cjk'),
    (r'\bconsequently\b', '从而', 'cjk'),
    (r'\boverlapping\b', '重叠', 'cjk'),
    (r'\bcontribut(?:e|ion)\b', '贡献', 'cjk'),
    (r'\bindependent(?:ly)?\b', '独立', 'cjk'),
    (r'\bimpossible\b', '不可能', 'cjk'),
    (r'\biterat(?:e|ion|ing)\b', '迭代', 'cjk'),
    (r'\benumerat(?:e|ion|ing)\b', '枚举', 'cjk'),

    # ── LOGIC SYMBOLS ──
    (r'\btherefore\b', '⇒', 'symbol'),
    (r'\bthus\b', '⇒', 'symbol'),
    (r'\bsuch that\b', 'st', 'abbrev'),
    (r'\bthere exists?\b', '∃', 'symbol'),
    (r'\bfor each\b', '∀', 'symbol'),
    (r'\bfor every\b', '∀', 'symbol'),
    (r'\bfor all\b', '∀', 'symbol'),
    (r'\bdoes not\b', '¬', 'symbol'),
    (r"\bdoesn't\b", '¬', 'symbol'),
    (r'\bis not\b(?!\s+(?:None|null|undefined|empty|zero|0))', '非', 'cjk'),  # protect "is not None" etc
    (r'\bat least\b', '≥', 'symbol'),
    (r'\bat most\b', '≤', 'symbol'),
    (r'\bgreater than\b', '>', 'symbol'),
    (r'\bless than\b', '<', 'symbol'),

    # ── ARTICLE STRIPPING (last — lowest priority) ──
    (r'\bthe\b\s+(?!(?:same|only|first|last|next|other)\b)', '', 'filler'),  # protect "the same", "the only" etc
    (r'\ba\b\s+(?=[bcdfghjklmnpqrstvwxyz])', '', 'filler'),
    (r'\ban\b\s+', '', 'filler'),
]

# ── Compile once ───────────────────────────────────────────────────────
_FILLER_COMPILED = [(re.compile(p, re.IGNORECASE), '') for p in _FILLER_PATTERNS]
_SUBS_COMPILED = [(re.compile(p, re.IGNORECASE), r, cat) for p, r, cat in _SUBSTITUTIONS]


def _ntok(text: str, tokenizer) -> int:
    """Token count using the provided tokenizer."""
    return len(tokenizer.encode(text, add_special_tokens=False))


def _protect_code_fences(text: str) -> tuple[str, list]:
    """Extract code-fenced blocks, replace with placeholders.
    Returns (text_with_placeholders, list_of_extracted_blocks)."""
    blocks = []
    def _replace(m):
        blocks.append(m.group(0))
        return f'\x00CODEFENCE{len(blocks)-1}\x00'
    # Match ```...``` and inline `...` (non-greedy)
    protected = re.sub(r'```.*?```|`[^`\n]+`', _replace, text, flags=re.DOTALL)
    return protected, blocks


def _restore_code_fences(text: str, blocks: list) -> str:
    """Restore code-fenced blocks from placeholders."""
    for i, block in enumerate(blocks):
        text = text.replace(f'\x00CODEFENCE{i}\x00', block)
    return text


def _apply_subs(text: str) -> tuple[str, dict]:
    """Apply all substitutions, return (result, stats).
    Code fences (``` and inline `) are protected from substitution."""
    stats = {'filler_drops': 0, 'cjk': 0, 'symbol': 0, 'abbrev': 0, 'total_subs': 0}

    # Phase 0: protect code fences from substitution
    text, code_blocks = _protect_code_fences(text)

    # Phase 1: filler drops
    for pat, repl in _FILLER_COMPILED:
        text, n = pat.subn(repl, text)
        if n:
            stats['filler_drops'] += n
            stats['total_subs'] += n

    # Phase 2: substitutions
    for pat, repl, cat in _SUBS_COMPILED:
        text, n = pat.subn(repl, text)
        if n:
            stats[cat] = stats.get(cat, 0) + n
            stats['total_subs'] += n

    # Phase 3: restore code fences
    text = _restore_code_fences(text, code_blocks)

    # Phase 4: whitespace normalization
    text = re.sub(r'[ \t]+', ' ', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = re.sub(r' *\n *', '\n', text)
    text = text.strip()

    return text, stats


def tokenmax(text: str, tokenizer, force_cjk: bool = False) -> str:
    """Apply token-maxing. Returns original if no savings achieved.

    Args:
        text: The think block content (without <think> tags).
        tokenizer: A HuggingFace tokenizer with .encode() method.
        force_cjk: If True, always return the processed version when CJK
            substitutions were applied, even if total token count increased.
            Use this to maximize CJK adoption in training data.

    Returns:
        The token-maxed text, or the original if processing didn't save tokens
        (unless force_cjk=True and CJK subs were applied).
    """
    if not text or not text.strip():
        return text

    original_tokens = _ntok(text, tokenizer)
    result, stats = _apply_subs(text)
    result_tokens = _ntok(result, tokenizer)

    # GUARD: only return processed version if it actually saves tokens
    # OVERRIDE: force_cjk bypasses the guard when CJK substitutions were made
    if result_tokens < original_tokens:
        return result
    if force_cjk and stats.get('cjk', 0) > 0:
        return result
    return text


def tokenmax_with_stats(text: str, tokenizer, force_cjk: bool = False) -> tuple[str, dict]:
    """Like tokenmax() but also returns substitution statistics.

    Args:
        force_cjk: If True, always apply when CJK substitutions were made,
            even if total token count increased. Prioritizes CJK adoption
            over token savings.

    Returns:
        (processed_text, stats_dict) where stats_dict contains:
          - original_tokens: token count before processing
          - result_tokens: token count after processing
          - saved: tokens saved (negative = token increase; check forced_cjk)
          - applied: whether the processed version was used
          - forced_cjk: True when force_cjk override caused acceptance despite no savings
          - filler_drops, cjk, symbol, abbrev: substitution counts by category
          - total_subs: total substitutions applied
    """
    if not text or not text.strip():
        return text, {'original_tokens': 0, 'result_tokens': 0, 'saved': 0,
                      'applied': False, 'forced_cjk': False, 'total_subs': 0}

    original_tokens = _ntok(text, tokenizer)
    result, stats = _apply_subs(text)
    result_tokens = _ntok(result, tokenizer)
    saved = original_tokens - result_tokens

    stats['original_tokens'] = original_tokens
    stats['result_tokens'] = result_tokens
    stats['saved'] = saved
    stats['forced_cjk'] = False

    if saved > 0:
        stats['applied'] = True
        return result, stats
    if force_cjk and stats.get('cjk', 0) > 0:
        stats['applied'] = True
        stats['forced_cjk'] = True
        return result, stats
    stats['applied'] = False
    return text, stats


# ── CLI: batch process a JSONL file ────────────────────────────────────
if __name__ == '__main__':
    import json, sys, argparse
    from transformers import AutoTokenizer

    parser = argparse.ArgumentParser(description='Token-max post-processor for think blocks')
    parser.add_argument('--input', required=True, help='Input JSONL (messages format)')
    parser.add_argument('--output', help='Output JSONL (default: dry run, stats only)')
    parser.add_argument('--tokenizer', default='ZelligeAI/tessera-compressor',
                        help='HF repo id or local path of the tokenizer to count savings under')
    parser.add_argument('--force-cjk', action='store_true',
                        help='Force CJK substitutions even if total tokens increase. '
                             'Prioritizes CJK adoption over token savings.')
    args = parser.parse_args()

    tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)

    total_before = total_after = applied = skipped = forced = 0

    out_lines = []
    with open(args.input) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            for m in rec.get('messages', rec.get('conversations', [])):
                role = m.get('role', m.get('from', ''))
                if role not in ('assistant', 'gpt'):
                    continue
                content_key = 'content' if 'content' in m else 'value'
                c = m.get(content_key, '') or ''
                if '<think>' not in c or '</think>' not in c:
                    continue

                # Extract think content, preserving prefix before <think> and suffix after </think>
                think_start = c.index('<think>') + len('<think>')
                think_end = c.index('</think>')
                prefix = c[:think_start - len('<think>')]
                think = c[think_start:think_end]
                suffix = c[think_end + len('</think>'):]

                maxed, stats = tokenmax_with_stats(think, tok, force_cjk=args.force_cjk)

                total_before += stats['original_tokens']
                if stats['applied']:
                    applied += 1
                    total_after += stats['result_tokens']
                    m[content_key] = f'{prefix}<think>{maxed}</think>{suffix}'
                    if stats.get('forced_cjk'):
                        forced += 1
                else:
                    skipped += 1
                    total_after += stats['original_tokens']

            out_lines.append(json.dumps(rec, ensure_ascii=False))

    if args.output:
        with open(args.output, 'w') as f:
            for line in out_lines:
                f.write(line + '\n')

    total = applied + skipped
    saved = total_before - total_after
    if total > 0:
        print(f'Processed {total} think blocks')
        print(f'  Applied: {applied} ({100*applied/total:.0f}%)')
        if forced:
            print(f'  Forced CJK: {forced} (applied despite no token savings)')
        print(f'  Skipped (no savings): {skipped}')
        pct = f'{100*saved/total_before:.1f}' if total_before > 0 else '0.0'
        print(f'  Tokens: {total_before} → {total_after} = {saved:+d} ({pct}%)')
    else:
        print(f'No think blocks found in {args.input}')