"""离线评测指标：纯函数实现，不依赖 Qdrant / Ollama，可独立单测

- entity_recall：实体保留率（摘要质量）
- routing_f1：L1 路由层 micro P/R/F1
- pruning_loss：上层剪枝损失率
- precision_at_k / recall_at_k：检索效用
- aggregate：按 key 汇总均值
"""

import re

# 数字串（含小数/版本号/百分数），如 1536、3.12、95%
_NUMBER_RE = re.compile(r"\d+(?:\.\d+)*%?")
# 型号/版本号模式，如 bge-m3、v1.2.0、FD-07
_MODEL_RE = re.compile(r"[A-Za-z]+[-_]\w*\d\w*|[A-Za-z]+\d+(?:\.\d+)*")
# 条款号，如 第三条、第 5 章、第十条
_CLAUSE_RE = re.compile(r"第\s*[一二三四五六七八九十百千万零0-9]+\s*[条款章节项]")
# 英文/大小写混合词（含连字符/点号连接），如 Qdrant、OpenAI、text-embedding-3-small
_EN_WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*(?:[-_.][A-Za-z0-9]+)*")
# 中文书名号内词，如 《部署手册》
_BOOK_TITLE_RE = re.compile(r"《([^》]+)》")


def _extract_entities(text: str) -> set[str]:
    """从文本抽取关键实体 token（去重集合）"""
    entities: set[str] = set()
    entities.update(_NUMBER_RE.findall(text))
    entities.update(_MODEL_RE.findall(text))
    entities.update(m.replace(" ", "") for m in _CLAUSE_RE.findall(text))
    entities.update(_BOOK_TITLE_RE.findall(text))
    # 英文词至少 2 个字符，避免单字母噪声
    entities.update(w for w in _EN_WORD_RE.findall(text) if len(w) >= 2)
    return entities


def entity_recall(source_text: str, summary: str) -> float:
    """实体保留率：原文抽取的关键实体在摘要中保留的比例

    原文无可抽取实体时返回 1.0（无实体可丢，视为满分）。
    """
    entities = _extract_entities(source_text)
    if not entities:
        return 1.0
    kept = sum(1 for entity in entities if entity in summary)
    return kept / len(entities)


def routing_f1(golden_doc_ids_per_query: list[set[str]], routed_doc_ids_per_query: list[set[str]]) -> dict[str, float]:
    """L1 路由层 micro precision / recall / f1

    逐 query 累加 tp/fp/fn 后统一计算（micro 平均）；
    两条列表必须等长，通常只统计 positive query（negative query 无 golden doc）。
    """
    tp = fp = fn = 0
    for golden, routed in zip(golden_doc_ids_per_query, routed_doc_ids_per_query, strict=True):
        tp += len(golden & routed)
        fp += len(routed - golden)
        fn += len(golden - routed)
    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
    return {"precision": precision, "recall": recall, "f1": f1}


def pruning_loss(cases: list[bool]) -> float:
    """剪枝损失率：原文含答案但 L1 未命中 golden doc 的 positive query 占比

    cases 中每个元素表示一条 positive query 是否被上层剪掉（True = 被剪）。
    """
    if not cases:
        return 0.0
    return sum(cases) / len(cases)


def precision_at_k(hit_doc_ids: list[str], golden: set[str], k: int) -> float:
    """Precision@k：前 k 条命中中 doc_id 属于 golden 的比例（按命中条数计，不去重）"""
    if k <= 0:
        return 0.0
    top = hit_doc_ids[:k]
    if not top:
        return 0.0
    hits = sum(1 for doc_id in top if doc_id in golden)
    return hits / k


def recall_at_k(hit_doc_ids: list[str], golden: set[str], k: int) -> float:
    """Recall@k：前 k 条命中覆盖的 golden doc 比例（按 doc 去重）；golden 为空返回 0.0"""
    if not golden:
        return 0.0
    found = {doc_id for doc_id in hit_doc_ids[:k] if doc_id in golden}
    return len(found) / len(golden)


def aggregate(records: list[dict[str, float]]) -> dict[str, float]:
    """按 key 汇总均值：对记录列表中每个指标求算术平均

    某 key 只在部分记录中出现时，按出现的记录取均值；空列表返回 {}。
    """
    if not records:
        return {}
    keys = list(dict.fromkeys(key for record in records for key in record))
    result: dict[str, float] = {}
    for key in keys:
        values = [record[key] for record in records if key in record]
        result[key] = sum(values) / len(values) if values else 0.0
    return result