#!/usr/bin/env python3
"""Cell 信噪比(SNR)分析：26 个 D 层 cell 里谁是混子？（纯离线，零 API）

信号 = 9 模型两两 JSD 均值 —— cell 把不同模型拉开的能力
噪声 = 模型内 bootstrap 对分 JSD —— 采样固有波动
SNR = 信号 / 噪声

三重验证：
  A. drop-one：去掉单 cell 后 9 模型 top-1 归因是否仍 9/9、最弱同族间距是否恶化
  B. SNR 前向贪心：按 SNR 降序加 cell，达成 9/9 的最小集，再后向修剪
  C. 后向消元：从全量 cell 起逐个剔除"删了不伤"的 cell
锚点：glm_53 两次独立全量样本的逐 cell test-retest JSD（真实噪声）
产出：/tmp/bfd/cell_snr_report.txt + /tmp/bfd/cell_snr.json
"""
import itertools
import json
import random
import sys
from collections import defaultdict

from pathlib import Path  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))  # 仓库根/site-packages, 使 evalharness 包可导入
from evalharness.fingerprint.engine import (build_d_normalized, compare_cells,  # noqa: E402
                    distributions_by_cell, jsd_bits, load_reference)

BFD = "/tmp/bfd"
R = str(Path(__file__).resolve().parent / 'references')

MODELS = [
    ("deepseek_v4_flash", "deepseek_v4_flash_fusion_reference.json", "DS"),
    ("deepseek_v4_flash_0731", "deepseek_v4_flash_0731_fusion_reference.json", "DS"),
    ("deepseek_v4_pro", "deepseek_v4_pro_fusion_reference.json", "DS"),
    ("glm_51", "glm_51_fusion_reference.json", "GLM"),
    ("glm_52", "glm52_vectron_fusion_reference.json", "GLM"),
    ("glm_53", "glm53_fusion_reference.json", "GLM"),
    ("kimi_k2_6", "kimi_k2_6_fusion_reference.json", "Kimi"),
    ("kimi_k2_7code", "kimi_k2_7code_fusion_reference.json", "Kimi"),
    ("kimi_k3", "kimi_k3_fusion_reference.json", "Kimi"),
]
NAMES = [m[0] for m in MODELS]
FAMILY = {m[0]: m[2] for m in MODELS}
WEAK_PAIRS = [("glm_51", "glm_52"), ("kimi_k2_6", "kimi_k2_7code")]

rng = random.Random(20260907)
LINES = []


def log(s=""):
    print(s)
    LINES.append(s)


def f3(x):
    return "  —  " if x is None else f"{x:.3f}"


# ---------- 载入 ----------
samples = {}
for n, _, _ in MODELS:
    with open(f"{BFD}/{n}/raw_answers.jsonl") as f:
        samples[n] = build_d_normalized([json.loads(l) for l in f])
with open(f"{BFD}/glm_53/rerun/raw_answers.jsonl") as f:
    rerun_samples = build_d_normalized([json.loads(l) for l in f])

dists = {n: distributions_by_cell(s) for n, s in samples.items()}
rerun_dist = distributions_by_cell(rerun_samples)
refs = {n: load_reference(f"{R}/{rf}")["cells"] for n, rf, _ in MODELS}

ALL_CELLS = sorted(set().union(*[set(r) for r in refs.values()])
                   | set().union(*[set(d) for d in dists.values()]))
log(f"cell 总数: {len(ALL_CELLS)}（电池口径 26 cell × 25 样本 = 650 D 请求）")
log()

# ---------- 有效性 / 类目数 ----------
valid_n = {}
for n in NAMES:
    cnt = defaultdict(int)
    for s in samples[n]:
        if s["cat"] == "valid" and s["norm"] is not None:
            cnt[s["cell"]] += 1
    for c in ALL_CELLS:
        valid_n[(n, c)] = cnt.get(c, 0)

cats = {}
for c in ALL_CELLS:
    seen = set()
    for n in NAMES:
        seen |= {s["norm"] for s in samples[n]
                 if s["cell"] == c and s["cat"] == "valid"}
    cats[c] = len(seen)

# ---------- 噪声：bootstrap 对分 ----------
def boot_noise(ans, rounds=30):
    n = len(ans)
    h = n // 2
    if h < 5:
        return None
    vals = []
    for _ in range(rounds):
        perm = list(ans)
        rng.shuffle(perm)
        a, b = defaultdict(int), defaultdict(int)
        for x in perm[:h]:
            a[x] += 1
        for x in perm[h:2 * h]:
            b[x] += 1
        vals.append(jsd_bits(dict(a), dict(b)))
    return sum(vals) / len(vals)


noise = {}
for c in ALL_CELLS:
    per = []
    for n in NAMES:
        ans = [s["norm"] for s in samples[n]
               if s["cell"] == c and s["cat"] == "valid"]
        j = boot_noise(ans)
        if j is not None:
            per.append(j)
    noise[c] = (sum(per) / len(per), len(per)) if per else (None, 0)

# ---------- 信号：跨模型两两 JSD ----------
signal = {}
for c in ALL_CELLS:
    js = []
    for a, b in itertools.combinations(NAMES, 2):
        da, db = dists[a].get(c), dists[b].get(c)
        if da and db and sum(da.values()) >= 10 and sum(db.values()) >= 10:
            js.append(jsd_bits(da, db))
    signal[c] = (sum(js) / len(js), len(js)) if js else (None, 0)

# ---------- SNR ----------
snr = {}
for c in ALL_CELLS:
    s, _ = signal[c]
    nz, _ = noise[c]
    if s is None or s < 1e-9:
        snr[c] = 0.0
    elif nz is None or nz < 1e-9:
        snr[c] = 99.0
    else:
        snr[c] = s / nz

# ---------- 预计算 (模型, 参考) 逐 cell JSD ----------
J = {}
for m in NAMES:
    for r in NAMES:
        entries, _ = compare_cells(dists[m], refs[r])
        J[(m, r)] = {e["cell"]: e["jsd"] for e in entries}


def state(active):
    res, margins = {}, []
    for m in NAMES:
        vals = {}
        for r in NAMES:
            js = [J[(m, r)][c] for c in active if c in J[(m, r)]]
            if js:
                vals[r] = sum(js) / len(js)
        if not vals:
            res[m] = (False, None, None)
            continue
        best = min(vals, key=vals.get)
        own = vals.get(m)
        sib = [v for r, v in vals.items() if FAMILY[r] == FAMILY[m] and r != m]
        margin = (min(sib) - own) if (sib and own is not None) else None
        res[m] = (best == m, own, margin)
        if margin is not None:
            margins.append((margin, m))
    correct = sum(1 for m in NAMES if res[m][0])
    weakest = min(margins) if margins else None
    return res, correct, weakest


base_res, base_correct, base_weak = state(ALL_CELLS)
log("【基线（全 cell）】top-1 正确 %d/9；各模型同族间距（正值=安全）:" % base_correct)
for m in NAMES:
    _, own, mg = base_res[m]
    log(f"  {m:24s} own={f3(own)}  同族间距={f3(mg)}")
log(f"  最弱环节: {base_weak[1]} 间距 {base_weak[0]:+.3f}")
log()

# ---------- drop-one ----------
drop1 = {}
for c in ALL_CELLS:
    _, cor, wk = state([x for x in ALL_CELLS if x != c])
    drop1[c] = (cor, wk)

# ---------- 后向消元 ----------
active = list(ALL_CELLS)
removed_order = []
while True:
    cands = []
    for c in active:
        rest = [x for x in active if x != c]
        _, cor, wk = state(rest)
        if cor == 9:
            cands.append((wk[0], c))
    if not cands:
        break
    cands.sort(reverse=True)
    pick = cands[0][1]
    removed_order.append(pick)
    active = [x for x in active if x != pick]
final_res, final_correct, final_weak = state(active)
log(f"【后向消元】可安全剔除 {len(removed_order)} 个，最小充分集 {len(active)} cell，"
    f"top-1 {final_correct}/9，最弱间距 {final_weak[1]} {final_weak[0]:+.3f}")
log(f"  剔除顺序: {', '.join(removed_order)}")
log(f"  保留集:   {', '.join(active)}")
log()

# ---------- SNR 前向贪心 ----------
order = sorted(ALL_CELLS, key=lambda c: (-(snr[c] if snr[c] < 90 else 99), c))
fw_active, fw_correct = [], 0
for c in order:
    fw_active.append(c)
    _, fw_correct, _ = state(fw_active)
    if fw_correct == 9:
        break
# 后向修剪
changed = True
while changed:
    changed = False
    for c in sorted(fw_active, key=lambda x: snr[x]):
        rest = [x for x in fw_active if x != c]
        _, cor, _ = state(rest)
        if cor == 9:
            fw_active = rest
            changed = True
            break
log(f"【SNR 前向贪心】最小集 {len(fw_active)} cell: {', '.join(fw_active)}")
log()

# ---------- 弱对载荷 cell ----------
log("【弱对载荷】版本级最难的两组，哪些 cell 在真正出力（两模型间 JSD top8）:")
for a, b in WEAK_PAIRS:
    per = []
    for c in ALL_CELLS:
        da, db = dists[a].get(c), dists[b].get(c)
        if da and db and sum(da.values()) >= 10 and sum(db.values()) >= 10:
            per.append((jsd_bits(da, db), c))
    per.sort(reverse=True)
    log(f"  {a} vs {b}: " + ", ".join(f"{c}({j:.3f})" for j, c in per[:8]))
log()

# ---------- 总表 ----------
log("【cell 信噪比排行】（按 SNR 降序；drop1=去掉该 cell 后 top-1 是否仍 9/9 + 最弱间距变化）")
log(f"{'cell':34s} {'valid均':>7s} {'<10数':>5s} {'类目':>4s} {'对数':>4s} "
    f"{'信号':>6s} {'噪声':>6s} {'SNR':>6s}  drop1")
ranked = sorted(ALL_CELLS, key=lambda c: -snr[c])
table_rows = []
for c in ranked:
    avg_v = sum(valid_n[(n, c)] for n in NAMES) / len(NAMES)
    below = sum(1 for n in NAMES if valid_n[(n, c)] < 10)
    s, pairs = signal[c]
    nz, nmod = noise[c]
    cor, wk = drop1[c]
    d1 = f"{'✓9/9' if cor == 9 else '✗破坏'} Δ{wk[0] - base_weak[0]:+.3f}" if wk else "?"
    tag = " ★核心" if c in active else (" ✂可剪" if cor == 9 else " ⚠载荷")
    log(f"{c:34s} {avg_v:7.1f} {below:5d} {cats[c]:4d} {pairs:4d} "
        f"{f3(s):>6s} {f3(nz):>6s} {snr[c]:6.2f}  {d1}{tag}")
    table_rows.append({"cell": c, "avg_valid": round(avg_v, 1), "below10": below,
                       "categories": cats[c], "pairs": pairs,
                       "signal": None if s is None else round(s, 4),
                       "noise": None if nz is None else round(nz, 4),
                       "snr": round(snr[c], 3), "drop1_correct": cor,
                       "drop1_min_margin": None if not wk else round(wk[0], 4),
                       "in_min_set": c in active})
log()

# ---------- test-retest 锚点 ----------
log("【真实噪声锚点】glm_53 两次独立全量样本的逐 cell test-retest JSD:")
tr = []
for c in ALL_CELLS:
    da, db = dists["glm_53"].get(c), rerun_dist.get(c)
    if da and db and sum(da.values()) >= 10 and sum(db.values()) >= 10:
        tr.append((jsd_bits(da, db), c))
tr.sort(reverse=True)
log("  " + ", ".join(f"{c}({j:.3f})" for j, c in tr))
glm_boot = [noise[c][0] for c in ALL_CELLS
            if noise[c][0] is not None
            and valid_n[("glm_53", c)] >= 10]
log(f"  均值 {sum(j for j, _ in tr) / len(tr):.3f}（同批 cell 的 bootstrap 噪声均值 "
    f"{sum(glm_boot) / len(glm_boot):.3f}，两者同量级则 bootstrap 估计可信）")
log()

# ---------- 节省估算 ----------
K = len(active)
log("【节省估算】(纯 D 层剪枝，文本层 36 + 基线 5 不变)")
log(f"  保留 {K} cell → D 请求 {650} → {K * 25}（{K * 25 / 691 * 100:.0f}% 原量）")
for nm, old_min, d_req in [("kimi_k3", 65, 5.6), ("glm_53", 23, 2.0), ("deepseek_v4_flash", 22, 2.0)]:
    d_time = old_min - 4  # 文本层+基线约 4 分钟
    new_min = d_time * (K * 25) / 650 + 4
    log(f"  {nm:22s} verify {old_min} 分钟 → 约 {new_min:.0f} 分钟")
log()

# ---------- 落盘 ----------
with open(f"{BFD}/cell_snr_report.txt", "w") as f:
    f.write("\n".join(LINES) + "\n")
with open(f"{BFD}/cell_snr.json", "w") as f:
    json.dump({"baseline": {"correct": base_correct,
                            "weak": [base_weak[1], round(base_weak[0], 4)]},
               "cells": table_rows,
               "backward_removed_order": removed_order,
               "minimal_set": active,
               "forward_set": fw_active,
               "final_margins": {m: None if base_res[m][2] is None else round(base_res[m][2], 4)
                                 for m in NAMES},
               "pruned_margins": {m: None if final_res[m][2] is None else round(final_res[m][2], 4)
                                  for m in NAMES}}, f, ensure_ascii=False, indent=1)
print("\n已写入 /tmp/bfd/cell_snr_report.txt + /tmp/bfd/cell_snr.json")