"""fake_embed_utils.py —— 自测用假嵌入（主题向量 + 词级 shingle 哈希）。

用途：在「不连生产 ollama/bge-m3」的纪律下，让 flow_docs 检索管道与评估集
可以在本地跑通全链路（切块 → 索引 → 检索 → 命中/引用断言）。

- 主题向量：把文本映射到粗主题组（identify/rnaseq/wgs/gwas/pop/motif/…），
  模拟 bge-m3 的语义召回
- 词级 shingle：英文词 ≥4 字符生成 4-char shingles、中文连续段生成 2-gram，
  哈希到向量维度，模拟词面重叠带来的排序区分度
- 权重：主题 0.35 + 词级 0.65，归一化后用于 IndexFlatIP（余弦）

注意：这是管道自测替身，不代表 bge-m3 真实质量；真实嵌入评测在部署环境用
run_flow_eval.py 默认模式（真 embed_batch）执行。
"""
from __future__ import annotations

import re
from typing import Iterable

import numpy as np

DIM = 1024
TOPIC_WEIGHT = 0.6
LEX_WEIGHT = 0.4
LEX_SAMPLE_CAP = 0     # 0 = 不采样（采样会丢失文档中部词；真实 bge-m3 全量编码）

_TOPIC_GROUPS = {
    "identify": ["identify", "初筛", "HMMER", "基因家族鉴定"],
    "motif": ["motif", "MEME", "基序"],
    "phylogeny": ["phylogeny", "系统发育", "MAFFT", "IQ-TREE", "建树"],
    "rnaseq": ["rnaseq", "RNA-seq", "DESeq2", "HISAT2", "转录组"],
    "wgs": ["wgs", "WGS", "BWA", "gvcf", "GATK", "重测序"],
    "gwas": ["gwas", "GWAS", "EMMAX", "kinship", "TASSEL", "关联"],
    "pop": ["pop", "群体", "PCA", "ADMIXTURE"],
    "prepare": ["genome_prepare", "基因组数据整理"],
    "promoter": ["promoter", "启动子", "PlantCARE"],
    "tfbs": ["tfbs", "转录因子结合位点", "FIMO"],
    "structure": ["structure", "外显子", "内含子", "基因结构"],
    "pep": ["pep", "ProtParam", "亚细胞定位", "跨膜"],
    "expression": ["expression", "表达模式", "TPM", "FPKM"],
    "report": ["report", "正式报告", "交付"],
    "delivery": ["交付目录", "项目交付", "00_input_and_design"],
    "synteny": ["synteny", "共线性"],
    "kaks": ["kaks", "Ka/Ks", "选择压力"],
    "ortholog": ["ortholog", "BLASTP", "直系同源"],
    "annotation": ["annotation", "gene2go", "功能注释", "InterProScan"],
    "de": ["rnaseq_de", "差异表达", "log2FoldChange", "padj"],
    "quant": ["rnaseq_quant", "定量", "featureCounts", "count_matrix"],
    "enrich": ["rnaseq_enrich", "富集", "clusterProfiler", "gene2ko"],
    "wgcna": ["wgcna", "WGCNA", "共表达", "软阈值", "soft-threshold"],
    "fgsea": ["fgsea", "gene2ko.txt", "DE_results.txt"],
    "candidate": ["candidate_family", "候选基因", "优先级"],
    "gvcf": ["gvcf", "HaplotypeCaller", "CombineGVCFs", "GenotypeGVCFs"],
    "impute": ["impute", "Beagle", "填补"],
    "filter": ["variant_filter", "硬过滤", "QD<2", "variant_select"],
    "markdup": ["markdup", "重复 reads"],
    "ld": ["ld_decay", "连锁不平衡", "LD"],
    "density": ["snp_density", "SNP 密度"],
    "cov": ["gwas_cov", "协变量"],
    "emmax": ["gwas_emmax", "EMMAX", "emmax_pvalues"],
    "glm": ["gwas_glm", "GLM"],
    "mlm": ["gwas_mlm", "MLM", "CMLM"],
    "checklist": ["质控清单", "专项检查"],
}

_CJK_RE = re.compile(r"[一-鿿]+")
_ASCII_WORD_RE = re.compile(r"[A-Za-z0-9_]+")


def _topic_of(text: str) -> str:
    for group, keys in _TOPIC_GROUPS.items():
        if any(k in text for k in keys):
            return group
    return "general"


def _lex_tokens(text: str) -> Iterable[str]:
    """词级 token：英文 4-char shingles + 中文 2-gram（长文本采样上限 LEX_SAMPLE_CAP）。"""
    lowered = text.lower()
    # 英文词（含数字/下划线）
    for m in _ASCII_WORD_RE.finditer(lowered):
        w = m.group(0)
        if len(w) >= 4:
            yield w                      # 整词
            for i in range(len(w) - 3):  # 4-gram shingles
                yield w[i:i + 4]
    # 中文连续段 2-gram（单字段取单字）
    for m in _CJK_RE.finditer(text):
        seg = m.group(0)
        if len(seg) == 1:
            yield seg
        else:
            for i in range(len(seg) - 1):
                yield seg[i:i + 2]


def _lex_tokens_sampled(text: str) -> list[str]:
    """采样版词级 token：上限 LEX_SAMPLE_CAP，保证长文本不占数量优势。"""
    toks = list(_lex_tokens(text))
    if len(toks) <= LEX_SAMPLE_CAP:
        return toks
    # 首尾各半采样：保留开头（标题/主题）与结尾（结论）分布
    half = LEX_SAMPLE_CAP // 2
    return toks[:half] + toks[-half:]


def fake_embed_flow(texts: list[str]) -> list[list[float]]:
    """主题 + 词级 shingle 混合假嵌入。空文本返回零向量。"""
    out: list[list[float]] = []
    for text in texts:
        if not text or not text.strip():
            out.append([0.0] * DIM)
            continue
        vec = np.zeros(DIM, dtype=np.float32)
        topic = _topic_of(text)
        vec[hash(topic) % DIM] += TOPIC_WEIGHT
        if topic == "general":
            # 无主题文本给平缓底噪，便于兜底召回
            vec[:16] += TOPIC_WEIGHT * 0.25
        n_tokens = 0
        toks = list(_lex_tokens(text))
        if LEX_SAMPLE_CAP > 0 and len(toks) > LEX_SAMPLE_CAP:
            half = LEX_SAMPLE_CAP // 2
            toks = toks[:half] + toks[-half:]
        for tok in toks:
            vec[hash(tok) % DIM] += LEX_WEIGHT
            n_tokens += 1
        if n_tokens:
            vec /= float(np.linalg.norm(vec))
        out.append(vec.tolist())
    return out