#!/usr/bin/env python3
"""FP-Fusion 维度 A：家族归因（Model Family Attribution）。

双路归因：
  1. 词表归因（主路，无需额外依赖）—— 按探针层加权命中家族别名：
       I 层自称命中      +1.5/次（含越狱/中英/直问）
       K 层组织/创建者   +1.0/次
       C 层拒答风格      +0.5/次（安全对齐措辞属家族化特征）
       S 层风格特征词    +0.3/次
       requested_family 先验 +0.8（served-name 声称家族）
  2. LLMmap 嵌入归因（辅路，可选）—— 对 I/K 层回答用 e5 embedding 与
     60 个已知模板算距离，取 top-3 模板的家族投票。

融合：
  S_fam = confidence × (1 - 0.3 × conflict)
  conflict = 主路 top1 家族与 LLMmap top1 家族不一致（量化/蒸馏/伪装信号）
"""

import os
import sys

from .scorer import _families_in_text, load_aliases


# -- 风格特征词（跨语言，家族化措辞弱信号）------------------------------
# 注意：这是"说话风格"侧写，不是身份声称，权重最低。
_STYLE_TOKENS = {
    # 英文：高频开场/缓冲词
    'certainly!', 'let me', 'let\'s', 'absolutely!', 'great question',
    'sure!', 'as an ai', 'i\'m happy to', 'i\'d be happy to',
    'good question', 'that\'s a great', 'here\'s how', 'let me break',
    # 中文
    '当然', '让我', '好的，', '没问题', '这是一个好问题',
    '首先', '其次', '最后', '总之',
}


def _lexicon_scores(records, aliases, requested_family):
    """词表归因：按层加权累加家族分数。records 需含 layer/response/error。"""
    scores = {}
    claim_hits = 0          # I 层可解析自称次数（归一化用）
    style_hits = 0

    for r in records:
        if r.get('error') or not r.get('response'):
            continue
        text = r.get('response') or ''
        layer = r.get('layer')
        fams = _families_in_text(text, aliases)

        if layer == 'I':
            if len(fams) == 1:                       # 明确自称
                f = next(iter(fams))
                scores[f] = scores.get(f, 0) + 1.5
                claim_hits += 1
            elif len(fams) > 1:                      # 多家族提及，分散加权
                for f in fams:
                    scores[f] = scores.get(f, 0) + 1.5 / len(fams)
                claim_hits += 1
        elif layer == 'K':
            for f in fams:                           # 组织/创建者弱证据
                scores[f] = scores.get(f, 0) + 1.0
        elif layer == 'C':
            low = text.lower()
            if any(s in low for s in ("i can't", "i cannot", "无法", "抱歉",
                                       "sorry", "i apologize")):
                for f in fams:                       # 拒答风格命中
                    scores[f] = scores.get(f, 0) + 0.5
                if not fams:
                    style_hits += 1                  # 拒答但未点名家族 → 中性
        elif layer == 'S':
            low = text.lower()
            for tok in _STYLE_TOKENS:
                if tok in low:
                    style_hits += 1
                    break
            # 风格词与家族弱相关：仅当该回答同时提及家族才累加
            if fams:
                for f in fams:
                    scores[f] = scores.get(f, 0) + 0.3

    # served-name 声称家族先验
    if requested_family and requested_family in aliases:
        scores[requested_family] = scores.get(requested_family, 0) + 0.8

    return scores, claim_hits, style_hits


def _normalize(scores):
    """按分数排序，返回 (top1_family, confidence, per_family)。"""
    if not scores:
        return None, 0.0, {}
    total = sum(scores.values())
    order = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)
    top1 = order[0][0]
    conf = order[0][1] / total if total > 0 else 0.0
    return top1, conf, dict(order)


def _llmmap_attribution(records, llmmap_tool):
    """LLMmap 嵌入归因（辅路）：I/K 层回答 → 模板距离 → 家族投票。

    llmmap_tool: 已加载的 LLMmap InferenceModel_open 实例（或 None）。
    返回 {top1_family, top3: [(family, dist)], votes: {family: n}}。
    """
    if llmmap_tool is None or not getattr(llmmap_tool, 'ready', False):
        return None
    import numpy as np

    texts = []
    for r in records:
        if r.get('layer') not in ('I', 'K'):
            continue
        if r.get('error') or not r.get('response'):
            continue
        texts.append((r['id'], r['response']))

    # LLMmap 模型一次只能吃固定 8 条 queries 的回答；这里每 8 条一批，
    # 逐条与模板库比对后按 label_map 归集家族。
    votes = {}
    dist_rows = []
    for batch_start in range(0, len(texts), 8):
        batch = texts[batch_start:batch_start + 8]
        answers = [t[1] for t in batch]
        # 不足 8 条时补空串（LLMmap __call__ 强校验数量）
        answers = answers + [''] * (8 - len(answers))
        try:
            dists = llmmap_tool(answers)
        except Exception:
            continue
        # open-set 距离：越小越好，取每条回答 top1 模板
        order_idx = int(np.argmin(dists))
        model_name = llmmap_tool.label_map[order_idx]
        votes[model_name] = votes.get(model_name, 0) + 1
        dist_rows.append((model_name, float(dists[order_idx])))

    if not votes:
        return None
    top3 = sorted(dist_rows, key=lambda x: x[1])[:3]
    top1_model = max(votes, key=votes.get)
    return {'top1_templates': top3, 'votes': votes,
            'top1_model': top1_model}


def family_attribution(records, aliases=None, requested_family=None,
                       llmmap_tool=None):
    """维度 A 主入口：双路归因融合。

    Args:
        records: fp_fusion 原始记录（含 layer/response/error/id）。
        aliases: 家族别名表 dict（None → 默认 family_aliases.json）。
        requested_family: served-name 声称家族（None → 自动从 model 解析）。
        llmmap_tool: 可选 LLMmap 实例。

    Returns dict（写入报告 signals.family）:
        enabled, method, top1_family, confidence, per_family_scores,
        claims, style_hits, llmmap(可选), conflict
    """
    aliases = aliases or load_aliases()

    scores, claims, style_hits = _lexicon_scores(records, aliases,
                                                 requested_family)
    top1, conf, per_fam = _normalize(scores)

    llm = _llmmap_attribution(records, llmmap_tool) if llmmap_tool else None
    conflict = False
    if llm and llm.get('top1_model'):
        # LLMmap 模板名 → 家族（'Qwen/Qwen2-7B-Instruct' → qwen）
        tmpl_fams = _families_in_text(llm['top1_model'], aliases)
        llm_fam = next(iter(tmpl_fams)) if len(tmpl_fams) == 1 else None
        if top1 and llm_fam and llm_fam != top1:
            conflict = True

    s_fam = conf * (1.0 - 0.3 * int(conflict)) if top1 else 0.0

    return {
        'enabled': True,
        'method': 'lexicon' + ('+llmmap' if llm else ''),
        'top1_family': top1,
        'confidence': round(conf, 4),
        'per_family_scores': {k: round(v, 3) for k, v in per_fam.items()},
        'claims_parsed': claims,
        'style_hits': style_hits,
        'llmmap': llm,
        'llmmap_family': llm_fam if llm else None,
        'conflict': conflict,
        's_fam': round(s_fam, 4),
    }