#!/usr/bin/env python3
"""FP-Fusion 维度 B：变体区分探针（V 层，Variant Probes）。

针对"同家族/同模型不同变体（量化档 / Instruct-SFT / DPO / 蒸馏）"的区分，
信号分两路：
  * 灰盒（logprobs）：engine 已采集 top_logprobs，此处探针保证 logprobs 有
    高信息量的回答可采样（确定性 / 高精度算术 / 长文本保真）。
  * 黑盒（文本）：靠回答本身的一致性 / 精度比对。

探针约定（对齐 engine._run_text_layer）：
  * layer = 'V'
  * meta 键仅允许：temperature / pair / lang / refusal_grad / len_ctrl /
                     role / expect_family（engine meta 白名单会透传）
  * v_determinism 探针自带 temperature=0.0（engine 支持 per-probe 温度后生效）
"""

# -- V 层灰盒探针（配合 --logprobs） ------------------------------------
V_GRAYBOX_PROBES = [
    # 高精度算术：量化模型在多位乘法/连续运算上精度掉得快
    {'id': 'v_arith_mul', 'layer': 'V', 'meta': {'temperature': 0.0},
     'text': ('Calculate exactly: 37 × 29 × 46 = ? '
              'Answer with only the final number, no explanation.')},
    {'id': 'v_arith_big', 'layer': 'V', 'meta': {'temperature': 0.0},
     'text': ('What is 862,391 + 459,782 + 123,456? '
              'Answer with only the final number.')},
    # 长文本保真：让模型重复长串，量化对复制保真更敏感
    {'id': 'v_repeat', 'layer': 'V', 'meta': {'temperature': 0.0, 'max_tokens': 256},
     'text': ('Repeat exactly the following string three times, separated by '
              'spaces, with no other text: "ai-2026-bf16-fp8-128ktokens"')},
    # 确定性采样两次：temp=0 下同 prompt 两次一致率 = 确定性的代理
    {'id': 'v_determinism_a', 'layer': 'V', 'meta': {'temperature': 0.0},
     'text': 'What is 9 × 8 + 7? Answer with only the number.'},
    {'id': 'v_determinism_b', 'layer': 'V', 'meta': {'temperature': 0.0},
     'text': 'What is 9 × 8 + 7? Answer with only the number.'},
]

# -- V 层黑盒探针（无需 logprobs） --------------------------------------
V_BLACKBOX_PROBES = [
    # 大数字逐位还原：数字 token 对量化最敏感
    {'id': 'v_fact_digit', 'layer': 'V', 'meta': {'temperature': 0.7},
     'text': ('What is the tens digit of 298473? '
              'Answer with only that single digit.')},
    # 指令精读：精确记忆并执行多约束指令
    {'id': 'v_instruction_exact', 'layer': 'V', 'meta': {'temperature': 0.7},
     'text': ('Do exactly 3 things in this order, each on its own line: '
              '1) write the number 7 2) write the word "blue" 3) write "done". '
              'Do not add anything else.')},
]


def ALL_VARIANT_PROBES():
    """V 层全量探针（variant / robustness 模式使用）。"""
    return V_GRAYBOX_PROBES + V_BLACKBOX_PROBES


def variant_signal(records, logprobs_enabled=False, notes=None):
    """维度 B 信号：从记录（含 top_logprobs）提炼变体区分指标。

    输出块（写入 report['signals']['variant']）：
      enabled, graybox_present, logprob_mean (top-1 logprob 均值),
      top1_stability (temp=0 同问两次 top1 一致率),
      self_consistency_jsd (黑盒：v_determinism a/b 回答一致率),
      arith_precision, notes
    """
    import re

    gray = [r for r in records if r.get('top_logprobs')]
    graybox_present = logprobs_enabled and len(gray) > 0

    # 1) 灰盒：top-1 logprob 均值 + 前 3 token 稳定率
    logprob_sum, logprob_cnt = 0.0, 0
    topk_stable = topk_total = 0
    for r in gray:
        for tk in (r.get('top_logprobs') or []):
            top = tk.get('top') or []
            if top:
                logprob_sum += top[0].get('logprob', 0.0)
                logprob_cnt += 1
    # top1 稳定性：v_determinism_a/b 同题两答的首 token top1 是否一致
    det_sigs = []
    for r in gray:
        if r.get('id') not in ('v_determinism_a', 'v_determinism_b'):
            continue
        tks = r.get('top_logprobs') or []
        det_sigs.append(tks[0]['top'][0]['token'] if tks and len(tks) > 0
                        and (tks[0].get('top') or []) else None)
    top1_stability = None
    if len(det_sigs) == 2 and det_sigs[0] is not None and det_sigs[1] is not None:
        top1_stability = 1.0 if det_sigs[0] == det_sigs[1] else 0.0

    # 2) 黑盒自一致性：v_determinism_a/b 回答相等
    det_a = next((r.get('response') or '').strip().lower()
                 for r in records if r.get('id') == 'v_determinism_a') if any(
                     r.get('id') == 'v_determinism_a' for r in records) else ''
    det_b = next((r.get('response') or '').strip().lower()
                 for r in records if r.get('id') == 'v_determinism_b') if any(
                     r.get('id') == 'v_determinism_b' for r in records) else ''
    det_match = bool(det_a and det_b and det_a == det_b)
    self_consistency_jsd = 1.0 if det_match else 0.0

    # 3) 算术精度：v_arith_mul 回答是否等于 37*29*46=49358
    arith_precision = None
    for r in records:
        if r.get('id') == 'v_arith_mul':
            ans = re.sub(r'[^0-9]', '', r.get('response') or '')
            arith_precision = (ans == '49358')
            break

    return {
        'enabled': True,
        'graybox_present': graybox_present,
        'graybox_records': len(gray),
        'logprob_mean': round(logprob_sum / logprob_cnt, 3) if logprob_cnt else None,
        'top1_stability': round(top1_stability, 3) if top1_stability is not None else None,
        'self_consistency_jsd': round(self_consistency_jsd, 3),
        'arith_precision': arith_precision,
        'notes': notes or [],
    }