"""
vocab.txt 覆盖率分析
对每个训练文件统计：
- token 级覆盖率：训练文件中的词有多少 % 在 vocab.txt 里
- 类型级覆盖率：训练文件中的唯一词有多少 % 在 vocab.txt 里
- OOV（out-of-vocab）最高频词 Top 20
"""

import re
from pathlib import Path
from collections import Counter

VOCAB_PATH  = Path(__file__).parent.parent / "evaluation-pipeline-2025/evaluation_pipeline/ewok/vocab.txt"
DATA_DIR    = Path(__file__).parent.parent / "data/raw_files"

FILES = [
    "childes.train.txt",
    "gutenberg.train.txt",
    "open_subtitles.train.txt",
    "simple_wiki.train.txt",
    "bnc_spoken.train.txt",
    "open_subtitles_cleaned_ckpt_70.txt",
    "switchboard.train.txt",
]

def tokenize(line):
    return re.findall(r"[a-zA-Z']+|[0-9]+", line.lower())

def analyze(file_name, vocab):
    path = DATA_DIR / file_name
    total_tokens = 0
    in_vocab_tokens = 0
    total_counter = Counter()
    oov_counter = Counter()

    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            tokens = tokenize(line)
            for t in tokens:
                total_tokens += 1
                total_counter[t] += 1
                if t in vocab:
                    in_vocab_tokens += 1
                else:
                    oov_counter[t] += 1

    total_types  = len(total_counter)
    in_vocab_types = sum(1 for w in total_counter if w in vocab)

    token_cov = in_vocab_tokens / total_tokens * 100 if total_tokens else 0
    type_cov  = in_vocab_types  / total_types  * 100 if total_types  else 0

    return {
        "file":           file_name,
        "total_tokens":   total_tokens,
        "in_vocab_tokens":in_vocab_tokens,
        "token_coverage": round(token_cov, 2),
        "total_types":    total_types,
        "in_vocab_types": in_vocab_types,
        "type_coverage":  round(type_cov, 2),
        "top20_oov":      oov_counter.most_common(20),
    }

def main():
    print(f"读取 vocab.txt ...")
    vocab = set(Path(VOCAB_PATH).read_text().splitlines())
    print(f"vocab 大小: {len(vocab):,} 词\n")

    results = []
    for fname in FILES:
        path = DATA_DIR / fname
        if not path.exists():
            print(f"[skip] {fname} 不存在")
            continue
        print(f"分析: {fname} ...")
        r = analyze(fname, vocab)
        results.append(r)
        print(f"  token覆盖率: {r['token_coverage']}%  |  类型覆盖率: {r['type_coverage']}%")
        print(f"  总词数: {r['total_tokens']:,}  vocab内: {r['in_vocab_tokens']:,}")
        print(f"  总唯一词: {r['total_types']:,}  vocab内: {r['in_vocab_types']:,}")
        print(f"  Top10 OOV: {[w for w,c in r['top20_oov'][:10]]}")
        print()

    # 汇总表
    print("\n" + "="*70)
    print(f"{'文件':<40} {'token覆盖率':>10} {'类型覆盖率':>10} {'总词数':>12}")
    print("="*70)
    for r in results:
        print(f"{r['file']:<40} {r['token_coverage']:>9}%  {r['type_coverage']:>9}%  {r['total_tokens']:>12,}")

    # 保存结果
    import json
    out = Path(__file__).parent.parent / "results/analysis/vocab_coverage.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(results, ensure_ascii=False, indent=2))
    print(f"\n结果已保存: {out}")

if __name__ == "__main__":
    main()