"""示例共用：确定性的“假嵌入”（随机向量 + 主题偏移），让示例不依赖嵌入模型 API。"""
import numpy as np

DIM = 64
TOPICS = ["hr", "it", "finance"]


def make_corpus(n_per_topic: int = 200, seed: int = 7):
    """每个主题一个中心向量，文档向量 = 中心 + 噪声。返回 (ids, vectors, metadatas)。"""
    rng = np.random.default_rng(seed)
    centers = {t: rng.normal(size=DIM) for t in TOPICS}
    ids, vecs, metas = [], [], []
    for t in TOPICS:
        for i in range(n_per_topic):
            v = centers[t] + rng.normal(scale=0.8, size=DIM)
            ids.append(f"{t}-{i}")
            vecs.append(v / np.linalg.norm(v))
            metas.append({"topic": t, "tenant": f"tenant-{i % 3}"})
    return ids, np.array(vecs, dtype=np.float32), metas, centers


def query_vec(centers, topic: str, seed: int = 1) -> np.ndarray:
    rng = np.random.default_rng(seed)
    v = centers[topic] + rng.normal(scale=0.8, size=DIM)
    return (v / np.linalg.norm(v)).astype(np.float32)


def brute_force(vecs: np.ndarray, q: np.ndarray, k: int, mask=None) -> list[int]:
    """精确的最近邻（余弦 = 归一化后的点积），作为近似检索的对照。"""
    scores = vecs @ q
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    return list(np.argsort(-scores)[:k])