"""
BLiMP + BLiMP Supplement 覆盖率分析

对每个子任务：
1. 提取 good 句中的"判别性关键词"（good有但bad没有的词）
2. 统计这些词在训练数据（3_dedup）里的总频次
3. 输出每个子任务的覆盖得分，找出弱点
"""

import json, re
from pathlib import Path
from collections import Counter, defaultdict

ROOT        = Path(__file__).parent.parent.parent
BLIMP_DIR   = ROOT / "evaluation-pipeline-2025/evaluation_data/full_eval/blimp_filtered"
SUPP_DIR    = ROOT / "evaluation-pipeline-2025/evaluation_data/full_eval/supplement_filtered"
TRAIN_DIR   = ROOT / "data/3_dedup"
OUT_DIR     = ROOT / "results/analysis"
OUT_DIR.mkdir(parents=True, exist_ok=True)

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

def tokenize(text):
    return re.findall(r"[a-z']+", text.lower())

# ── 1. 统计训练数据词频 ───────────────────────────────────
print("统计训练数据词频...")
train_freq = Counter()
train_total = 0
for fname in TRAIN_FILES:
    path = TRAIN_DIR / fname
    if not path.exists():
        continue
    with open(path, errors="ignore") as f:
        for line in f:
            words = tokenize(line)
            train_total += len(words)
            train_freq.update(words)
print(f"  训练数据总词数: {train_total:,}")

# ── 2. 分析每个BLiMP子任务 ───────────────────────────────
def analyze_task(jfile):
    good_words = Counter()
    bad_words  = Counter()
    n_pairs = 0
    with open(jfile) as f:
        for line in f:
            d = json.loads(line)
            good_words.update(tokenize(d["sentence_good"]))
            bad_words.update(tokenize(d["sentence_bad"]))
            n_pairs += 1

    # 判别性关键词：good句里有，bad句里没有（或少很多）
    discriminative = {}
    for w, gc in good_words.items():
        bc = bad_words.get(w, 0)
        if gc - bc >= max(3, gc * 0.3):   # good比bad多30%以上且至少差3次
            discriminative[w] = gc - bc

    # 这些词在训练数据里的频次（每百万词）
    if not discriminative:
        return None

    top_words = sorted(discriminative.items(), key=lambda x: -x[1])[:10]
    coverage_scores = []
    for w, diff in top_words:
        per_million = train_freq.get(w, 0) / train_total * 1e6
        coverage_scores.append((w, diff, per_million))

    # 子任务覆盖得分 = 关键词在训练数据中的平均频次（每百万词）
    avg_coverage = sum(s for _, _, s in coverage_scores) / len(coverage_scores)

    return {
        "task":         jfile.stem,
        "n_pairs":      n_pairs,
        "top_keywords": coverage_scores,
        "avg_coverage": avg_coverage,
    }

# 分析BLiMP
print("\n分析 BLiMP 子任务...")
blimp_results = []
for jfile in sorted(BLIMP_DIR.glob("*.jsonl")):
    r = analyze_task(jfile)
    if r:
        blimp_results.append(r)

# 分析BLiMP Supplement
print("分析 BLiMP Supplement 子任务...")
supp_results = []
for jfile in sorted(SUPP_DIR.glob("*.jsonl")):
    r = analyze_task(jfile)
    if r:
        supp_results.append(r)

# ── 3. 输出结果 ──────────────────────────────────────────
def print_results(results, title):
    print(f"\n{'='*70}")
    print(f"{title}（按覆盖率从低到高，最需要补充的在前）")
    print(f"{'='*70}")
    print(f"{'子任务':<45} {'覆盖分':>8}  {'关键词(每百万词)'}")
    print("─" * 70)
    for r in sorted(results, key=lambda x: x["avg_coverage"]):
        kw_str = ", ".join(f"{w}({s:.0f})" for w, _, s in r["top_keywords"][:4])
        print(f"{r['task']:<45} {r['avg_coverage']:>8.1f}  {kw_str}")

print_results(blimp_results,  "BLiMP 子任务覆盖分析")
print_results(supp_results,   "BLiMP Supplement 覆盖分析")

# 保存完整结果
import json as _json
out_path = OUT_DIR / "blimp_coverage.json"
_json.dump({"blimp": blimp_results, "supplement": supp_results},
           open(out_path, "w"), ensure_ascii=False, indent=2)
print(f"\n完整结果已保存: {out_path}")