#!/usr/bin/env python3
"""剪枝候选集验证（纯离线，零 API）：

确定性检查之外，加 bootstrap 重跑仿真：对每个模型的 cell 答案做有放回重采样，
模拟"明天再跑一遍电池"，统计 top-1 归因成功率。剪枝集与全量集成功率持平才算安全。

集合定义（依据 cell_snr.py 排行 + 弱对载荷分析）：
  tier1 死重: binary-*-zh ×4（9 模型全 0 valid，纯浪费 100 请求/跑）+ favorite-number:en（5/9 模型不可用）
  tier2 零信号: coin-flip:en/zh、random-number-1-10:en/zh、binary-season:en
               （跨模型信号 ≤0.053，全体模型收敛到同一分布，构造性无区分力）
  tier3 谨慎: day-of-week:zh（不在任何弱对 top8，移除 Δ+0.001）
  推荐15 = 弱对 top8 并集(11) + SNR≥2 补充(4)
"""
import json
import random
import sys
from collections import Counter, 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, distributions_by_cell,  # noqa: E402
                    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}

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])
dists = {n: distributions_by_cell(s) for n, s in samples.items()}
refs = {n: load_reference(f"{R}/{rf}")["cells"] for n, rf, _ in MODELS}

ALL = sorted(set().union(*[set(r) for r in refs.values()])
             | set().union(*[set(d) for d in dists.values()]))

ANS = {}
for n in NAMES:
    per = defaultdict(list)
    for s in samples[n]:
        if s["cat"] == "valid" and s["norm"] is not None:
            per[s["cell"]].append(s["norm"])
    for c in ALL:
        ANS[(n, c)] = per.get(c, [])


def score(dist_m, ref_cells, cells):
    js = []
    for c in cells:
        da, db = dist_m.get(c), ref_cells.get(c)
        if not da or not db:
            continue
        if sum(da.values()) < 10 or sum(db.values()) < 10:
            continue
        js.append(jsd_bits(da, db))
    return sum(js) / len(js) if js else None


def state(cells, D=None):
    D = dists if D is None else D
    res = {}
    for m in NAMES:
        vals = {r: score(D[m], refs[r], cells) for r in NAMES}
        vals = {r: v for r, v in vals.items() if v is not None}
        if not vals:
            res[m] = (None, False, False, 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]
        mg = (min(sib) - own) if (sib and own is not None) else None
        res[m] = (best, best == m, FAMILY[best] == FAMILY[m], mg)
    exact = sum(1 for m in NAMES if res[m][1])
    margins = [(res[m][3], m) for m in NAMES if res[m][3] is not None]
    return res, exact, (min(margins) if margins else None)


KEEP15 = [
    "random-animal:en", "random-animal:zh", "random-city:en", "random-city:zh",
    "random-color:en", "random-color:zh", "random-letter:en", "random-letter:zh",
    "random-number-1-100:en", "random-number-1-100:zh", "favorite-number:zh",
    "day-of-week:en", "binary-pet:en", "binary-tea-coffee:en",
    "binary-sea-mountain:en",
]
TIER12_16 = KEEP15 + ["day-of-week:zh"]
MINI3 = ["day-of-week:en", "binary-pet:en", "binary-tea-coffee:en"]

SETS = [
    ("全量 26 cell（现状）", ALL),
    ("推荐 15 cell", KEEP15),
    ("仅 tier1+2 剪（16 cell）", TIER12_16),
    ("家族分诊 mini 3 cell", MINI3),
]

out = []
def log(s=""):
    print(s)
    out.append(s)

B = 40
for name, cells in SETS:
    res, exact, weak = state(cells)
    d_req = len(cells) * 25
    k3_min = (65 - 4) * d_req / 650 + 4
    log(f"━━ {name}  D请求 {d_req}（{d_req / 691 * 100:.0f}% 原量）  K3 verify 约 {k3_min:.0f} 分钟")
    log(f"  确定性: 精确 top-1 {exact}/9  最弱间距 {weak[1]} {weak[0]:+.3f}")
    rng = random.Random(7)
    ex = {m: 0 for m in NAMES}
    fam = {m: 0 for m in NAMES}
    confusions = Counter()
    for _ in range(B):
        Dboot = {}
        for m in NAMES:
            d = {}
            for c in cells:
                ans = ANS[(m, c)]
                if len(ans) >= 10:
                    d[c] = dict(Counter(rng.choices(ans, k=len(ans))))
            Dboot[m] = d
        for m in NAMES:
            vals = {r: score(Dboot[m], refs[r], cells) for r in NAMES}
            vals = {r: v for r, v in vals.items() if v is not None}
            if not vals:
                continue
            best = min(vals, key=vals.get)
            ex[m] += best == m
            fam[m] += FAMILY[best] == FAMILY[m]
            if best != m:
                confusions[(m, best)] += 1
    tot_ex = sum(ex.values()) / (B * len(NAMES)) * 100
    tot_fam = sum(fam.values()) / (B * len(NAMES)) * 100
    log(f"  bootstrap 重跑仿真({B}次): 精确 {tot_ex:.0f}%  家族 {tot_fam:.0f}%")
    detail = "  ".join(f"{m.replace('deepseek_v4_', 'ds_').replace('kimi_', 'k')}: {ex[m] / B * 100:.0f}%"
                       for m in NAMES)
    log(f"    逐模型精确: {detail}")
    if confusions:
        top = ", ".join(f"{a}→{b}×{c}" for (a, b), c in confusions.most_common(4))
        log(f"    混淆集中在: {top}")
    log()

with open(f"{BFD}/keepset_eval.txt", "w") as f:
    f.write("\n".join(out) + "\n")
print("已写入 /tmp/bfd/keepset_eval.txt")