"""
训练 Morfessor+BPE Tokenizer

思路：
  1. 在训练语料上训练 Morfessor，学习英语形态素边界
  2. 用 Morfessor 对训练文本预切分（在形态素边界插入空格）
  3. 在预切分文本上训练 BPE（架构与标准 BPE tokenizer 完全一致）

效果：BPE 不会跨形态素边界合并，保留 WUG/Entity Tracking 有利的形态信息

输入: data/8_sample_B/train.txt
输出: models/tokenizer_morfessor/
      ├── tokenizer.json            (BPE tokenizer, HuggingFace 格式)
      ├── tokenizer_config.json
      ├── special_tokens_map.json
      └── morfessor.bin             (Morfessor 模型，推理时需要)

用法:
    pip install morfessor
    python scripts/02_model/train_tokenizer_morfessor.py
    python scripts/02_model/train_tokenizer_morfessor.py --input data/8_sample_C/train.txt
"""

import argparse
import re
import sys
from collections import Counter
from pathlib import Path

try:
    import morfessor
except ImportError:
    print("请先安装: pip install morfessor")
    sys.exit(1)

from tokenizers import Tokenizer, Regex
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.normalizers import Sequence, Prepend, NFKC, Replace
from tokenizers.pre_tokenizers import Sequence as PreSeq, Split, ByteLevel
from tokenizers.processors import TemplateProcessing
from transformers import PreTrainedTokenizerFast

ROOT = Path(__file__).resolve().parent.parent.parent
DEFAULT_INPUT = ROOT / "data/8_sample_B/train.txt"
DEFAULT_OUT = ROOT / "models/tokenizer_morfessor"
SPECIAL_TOKENS = ["<unk>", "<s>", "</s>", "<pad>", "<mask>"]

# Morfessor 超参
CORPUSWEIGHT = 0.01       # 越小切分越激进（默认~1.0 太保守）
MIN_MORPH_LEN = 2         # 形态素最短长度，过滤 s+it / b+and 等假阳性
MIN_WORD_LEN = 3          # 短于此长度的词不切分


# ══════════════════════════════════════════════════════════════════════════════
# Step 1: 训练 Morfessor
# ══════════════════════════════════════════════════════════════════════════════

def train_morfessor_model(input_path: Path, corpusweight: float = CORPUSWEIGHT) -> morfessor.BaselineModel:
    """在训练语料上训练 Morfessor 模型。"""
    print(f"Step 1: 训练 Morfessor 模型 (corpusweight={corpusweight})...")

    # 统计词频（只取纯字母词，去标点）
    word_counts = Counter()
    with open(input_path) as f:
        for line in f:
            for word in line.split():
                clean = word.strip(".,!?;:\"'()-[]{}…""''").lower()
                if len(clean) >= 2 and clean.isalpha():
                    word_counts[clean] += 1

    print(f"  唯一词数: {len(word_counts):,}")
    print(f"  总词频:   {sum(word_counts.values()):,}")

    model = morfessor.BaselineModel(corpusweight=corpusweight)
    training_data = [(count, word) for word, count in word_counts.items()]
    model.load_data(training_data)
    model.train_batch()

    # 示例
    samples = [
        "unhappiness", "running", "walked", "beautiful", "government",
        "internationally", "darkness", "singer", "swimming", "nationalization",
        "happiness", "slowly", "governmental", "children", "quickly",
    ]
    print("\n  切分示例:")
    for w in samples:
        segs = model.viterbi_segment(w)[0]
        print(f"    {w:25s} → {' + '.join(segs)}")

    return model


# ══════════════════════════════════════════════════════════════════════════════
# Step 2: 预切分文本
# ══════════════════════════════════════════════════════════════════════════════

_WORD_RE = re.compile(r'^([^a-zA-Z]*?)([a-zA-Z]+)([^a-zA-Z]*)$')


def presegment_word(word: str, morf_model) -> str:
    """对单词用 Morfessor 找形态素边界，在边界处插入空格。

    保留原始大小写：用 Morfessor 切分小写版本，
    再按 segment 长度映射回原始字符。

    过滤规则：所有形态素长度必须 >= MIN_MORPH_LEN，
    否则视为假阳性（如 s+it, b+and），保持原词不切。
    """
    m = _WORD_RE.match(word)
    if not m:
        return word

    prefix, core, suffix = m.groups()
    if len(core) < MIN_WORD_LEN:
        return word

    segments = morf_model.viterbi_segment(core.lower())[0]
    if len(segments) <= 1:
        return word

    # 过滤：所有形态素必须 >= MIN_MORPH_LEN
    if not all(len(s) >= MIN_MORPH_LEN for s in segments):
        return word

    # 按 segment 长度从原始 core 中切片（保留大小写）
    parts = []
    pos = 0
    for seg in segments:
        n = len(seg)
        parts.append(core[pos:pos + n])
        pos += n

    return prefix + ' '.join(parts) + suffix


def presegment_file(input_path: Path, morf_model, output_path: Path) -> Path:
    """对整个文件进行 Morfessor 预切分。"""
    print("\nStep 2: 预切分训练文本...")

    line_count = 0
    with open(input_path) as fin, open(output_path, 'w') as fout:
        for line in fin:
            if line.strip():
                words = line.split()
                segmented = [presegment_word(w, morf_model) for w in words]
                fout.write(' '.join(segmented) + '\n')
            else:
                fout.write('\n')
            line_count += 1
            if line_count % 200000 == 0:
                print(f"  已处理 {line_count:,} 行...")

    print(f"  总共处理 {line_count:,} 行")
    print(f"  预切分文件: {output_path} ({output_path.stat().st_size / 1e6:.1f} MB)")

    # 对比几行
    print("\n  对比示例:")
    shown = 0
    with open(input_path) as f1, open(output_path) as f2:
        for orig, seg in zip(f1, f2):
            if orig.strip() != seg.strip() and 30 < len(orig.strip()) < 150:
                print(f"    原: {orig.strip()}")
                print(f"    切: {seg.strip()}")
                print()
                shown += 1
                if shown >= 5:
                    break

    return output_path


# ══════════════════════════════════════════════════════════════════════════════
# Step 3: 在预切分文本上训练 BPE
# ══════════════════════════════════════════════════════════════════════════════

def build_and_train_bpe(segmented_file: Path, vocab_size: int) -> Tokenizer:
    """训练 BPE，架构与标准 tokenizer 完全一致。"""
    print(f"\nStep 3: 训练 BPE (vocab_size={vocab_size})...")

    # Normalizer（与标准 BPE 一致）
    normalizer = Sequence([
        Prepend(prepend=" "),
        NFKC(),
        Replace(Regex(r"\n"), "\n "),
        Replace(Regex(r" *\n"), "\n"),
    ])

    # Pre-tokenizer（与标准 BPE 一致）
    GPT4_REGEX = (
        r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*"
        r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+"
        r"|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+"
        r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]*"
        r"| ?\p{N}"
        r"| ?[^\s\p{L}\p{N}]+[\r\n/]*"
        r"|\s*[\r\n]+"
        r"|\s+(?!\S)"
        r"|\s+"
    )
    pre_tokenizer = PreSeq([
        Split(pattern=Regex(GPT4_REGEX), behavior="isolated"),
        ByteLevel(add_prefix_space=False, trim_offsets=True, use_regex=False),
        Split(pattern=Regex(r".{1,24}"), behavior="isolated"),
    ])

    tokenizer = Tokenizer(BPE(unk_token="<unk>"))
    tokenizer.normalizer = normalizer
    tokenizer.pre_tokenizer = pre_tokenizer

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=SPECIAL_TOKENS,
        min_frequency=2,
        show_progress=True,
    )

    tokenizer.train(files=[str(segmented_file)], trainer=trainer)
    print(f"  实际 vocab size: {tokenizer.get_vocab_size()}")

    # Post-processor: 句首加 <s>
    tokenizer.post_processor = TemplateProcessing(
        single="<s> $A",
        pair="<s> $A <s> $B",
        special_tokens=[("<s>", tokenizer.token_to_id("<s>"))],
    )

    # 校验特殊 token ID
    expected = {"<unk>": 0, "<s>": 1, "</s>": 2, "<pad>": 3, "<mask>": 4}
    for token, eid in expected.items():
        aid = tokenizer.token_to_id(token)
        status = "✅" if aid == eid else "❌"
        print(f"  {status}  {token:10s}  expected={eid}  actual={aid}")
        if aid != eid:
            raise ValueError(f"特殊 token ID 不一致: {token}")

    return tokenizer


# ══════════════════════════════════════════════════════════════════════════════
# Step 4: 保存
# ══════════════════════════════════════════════════════════════════════════════

def save_all(tokenizer: Tokenizer, morf_model, output_dir: Path):
    """保存 BPE tokenizer + Morfessor 模型。"""
    print(f"\nStep 4: 保存到 {output_dir}/")
    output_dir.mkdir(parents=True, exist_ok=True)

    # BPE tokenizer
    raw_path = output_dir / "tokenizer.json"
    tokenizer.save(str(raw_path))

    fast_tok = PreTrainedTokenizerFast(
        tokenizer_file=str(raw_path),
        bos_token="<s>", eos_token="</s>", unk_token="<unk>",
        sep_token="</s>", pad_token="<pad>", cls_token="<s>", mask_token="<mask>",
    )
    fast_tok.save_pretrained(str(output_dir))

    # Morfessor 模型
    io = morfessor.MorfessorIO()
    morf_path = output_dir / "morfessor.bin"
    io.write_binary_model_file(str(morf_path), morf_model)

    print(f"  文件: {sorted(f.name for f in output_dir.iterdir())}")


# ══════════════════════════════════════════════════════════════════════════════
# Step 5: 验证 & 对比
# ══════════════════════════════════════════════════════════════════════════════

def verify(output_dir: Path, morf_model):
    """加载并对比标准 BPE vs Morfessor+BPE。"""
    print("\nStep 5: 验证 & 对比...")

    fast_tok = PreTrainedTokenizerFast.from_pretrained(str(output_dir))

    std_path = ROOT / "models/tokenizer/tokenizer.json"
    if std_path.exists():
        std_tok = Tokenizer.from_file(str(std_path))
    else:
        std_tok = None

    tests = [
        "The cat sat on the mat.",
        "She was running quickly through the forest.",
        "I don't think he's coming today.",
        "unhappiness",
        "running jumped swimming",
        "The ice is cold and the fire is hot.",
        "nationalization",
        "governmental",
        "The children played happily in the garden.",
    ]

    for t in tests:
        # Morfessor+BPE: 先预切分
        seg_t = ' '.join(presegment_word(w, morf_model) for w in t.split())
        morf_tokens = fast_tok.tokenize(seg_t)

        print(f"  原文: {t}")
        if seg_t != t:
            print(f"  预切: {seg_t}")
        if std_tok:
            std_tokens = std_tok.encode(t).tokens
            print(f"    标准BPE   ({len(std_tokens):2d}): {std_tokens}")
        print(f"    Morf+BPE  ({len(morf_tokens):2d}): {morf_tokens}")
        print()


# ══════════════════════════════════════════════════════════════════════════════
# Main
# ══════════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="训练 Morfessor+BPE Tokenizer")
    parser.add_argument("--input", default=str(DEFAULT_INPUT), help="训练文件路径")
    parser.add_argument("--output", default=str(DEFAULT_OUT), help="输出目录")
    parser.add_argument("--vocab_size", default=8192, type=int, help="词表大小")
    parser.add_argument("--corpusweight", default=CORPUSWEIGHT, type=float,
                        help=f"Morfessor corpusweight, 越小越激进 (默认{CORPUSWEIGHT})")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_dir = Path(args.output)

    if not input_path.exists():
        raise FileNotFoundError(f"训练文件不存在: {input_path}")

    print(f"═══ 训练 Morfessor+BPE Tokenizer ═══")
    print(f"  输入:       {input_path} ({input_path.stat().st_size / 1e6:.1f} MB)")
    print(f"  输出:       {output_dir}")
    print(f"  vocab_size: {args.vocab_size}")
    print()

    # Step 1: Morfessor
    morf_model = train_morfessor_model(input_path, corpusweight=args.corpusweight)

    # Step 2: 预切分
    seg_path = output_dir / "_presegmented_train.txt"
    output_dir.mkdir(parents=True, exist_ok=True)
    presegment_file(input_path, morf_model, seg_path)

    # Step 3: BPE
    tokenizer = build_and_train_bpe(seg_path, args.vocab_size)

    # Step 4: 保存
    save_all(tokenizer, morf_model, output_dir)

    # Step 5: 验证
    verify(output_dir, morf_model)

    # 删除预切分中间文件
    seg_path.unlink()
    print(f"  已删除中间文件: {seg_path.name}")

    print("\n═══ 完成！═══")


if __name__ == "__main__":
    main()