# 导入默认字典 defaultdict，便于自动初始化桶结构
from collections import defaultdict

# 导入数学库，提供 sqrt、pow 等基础数学函数
import math

# 导入随机库，支持随机抽样与随机数生成
import random


# 计算两个向量之间的欧氏距离（L2 距离）
def l2_distance(a, b):
    # zip 将两个向量按维度配对，平方差求和后开平方
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


# K-Means++ 初始化：挑选 k 个初始中心，并做一次分配+更新
def kmeans(vectors, k):
    # 第一个中心直接从所有样本中随机挑选
    centers = [random.choice(vectors)]

    # 解释：下面这段 while 是 K-Means++ 的“轮盘赌”抽样逻辑
    # 步骤：计算每个点到最近中心的距离平方 -> 得到概率权重 -> 依概率抽样新的中心

    # 当中心数量不足 k 个时，持续抽样新的中心
    while len(centers) < k:
        # 第一步：每个数据点距离最近中心的平方，作为该点的权重
        dists = [min(l2_distance(v, c) ** 2 for c in centers) for v in vectors]
        # 第二步：对所有权重求和，得到轮盘赌的总长度
        total = sum(dists)

        # 第三步：在 [0, total) 区间上随机抽一个阈值 r，相当于轮盘指针
        r = random.random() * total

        # prefix 累加当前遍历到的权重，用于定位 r 落在哪个区间
        prefix = 0.0

        # 遍历所有点和对应权重，寻找累积首次超过 r 的那个点
        for v, d in zip(vectors, dists):
            # 不断累加权重，等价于在“条带”上向前走
            prefix += d
            # 一旦累计值覆盖了随机阈值 r，就把该点选为新中心
            if prefix >= r:
                centers.append(v)
                break

    # 初始化一个“簇字典”，key 是中心索引，value 是该中心下的样本列表
    clusters = {i: [] for i in range(k)}

    # 遍历所有向量，找到距离最近的中心，将该向量归入对应簇
    for v in vectors:
        idx = min(range(k), key=lambda i: l2_distance(v, centers[i]))
        clusters[idx].append(v)
    # 对每个非空簇计算均值，作为中心的最新位置
    for i, points in clusters.items():
        if points:
            # zip(*points) 将点按维度拆开，sum(dim)/len(points) 是维度均值
            centers[i] = [sum(dim) / len(points) for dim in zip(*points)]
    # 返回完成一次“初始化+更新”的中心列表
    return centers


# 构建 IVF-Flat 索引：输出中心列表和倒排桶
def build_ivf_flat(vectors, nlist):
    # 先通过 K-Means++ 得到 nlist 个中心
    centers = kmeans(vectors, nlist)

    # defaultdict(list) 便于向桶里 append，无需检查 key 是否存在
    buckets = defaultdict(list)

    # 遍历所有向量，依据最近中心编号将其放入对应桶
    for vec in vectors:
        idx = min(range(nlist), key=lambda i: l2_distance(vec, centers[i]))
        buckets[idx].append(vec)

    # 返回构建好的中心列表与倒排结构
    return centers, buckets


# 在 IVF-Flat 索引上执行查询，近似返回 topk 个最近向量
def search_ivf_flat(query, centers, buckets, nprobe, topk):
    # 先按查询向量到各中心的距离排序，得到“探查顺序”
    probe_order = sorted(
        range(len(centers)), key=lambda i: l2_distance(query, centers[i])
    )
    print(probe_order)
    print(buckets)

    # 用列表保存候选 (距离, 向量) 对
    candidate_pairs = []

    # 只深入距离最近的 nprobe 个中心对应的桶
    for idx in probe_order[:nprobe]:
        for vec in buckets[idx]:
            candidate_pairs.append((l2_distance(query, vec), vec))

    # 候选列表按距离升序排列，便于截取前 topk 个
    candidate_pairs.sort(key=lambda x: x[0])

    print(candidate_pairs)

    # 返回最接近查询的 topk 个候选
    return candidate_pairs[:topk]


# 主程序：演示 IVF-Flat 的构建与查询流程
def main():
    # 设定随机种子，保证示例结果可复现
    random.seed(0)

    # 构造 8 个二维向量作为“数据库”
    database = [
        [0.1, 0.2],
        [0.15, 0.18],
        [0.8, 0.85],
        [0.82, 0.79],
        [0.4, 0.4],
        [0.42, 0.45],
        [0.9, 0.1],
        [0.88, 0.15],
    ]

    # 查询向量：我们要找与它相似的数据库向量
    query = [0.12, 0.22]

    # nlist 表示建索引时划分的中心（簇）个数
    nlist = 4

    # nprobe 表示查询时会探查多少个最近的中心
    nprobe = 2

    # 构建 IVF-Flat 索引，得到中心与倒排桶
    centers, buckets = build_ivf_flat(database, nlist=nlist)

    # 执行检索，返回距离最近的 top3 个结果
    results = search_ivf_flat(query, centers, buckets, nprobe=nprobe, topk=3)

    # 打印本次检索的参数设置
    print(f"nlist={nlist}, nprobe={nprobe}")

    # 打印检索到的候选向量及其距离
    print("最近向量：")
    for dist, vec in results:
        print(f"向量={vec}, 距离={dist:.4f}")


# Python 入口：若当前脚本以 main 方式执行，则运行示例
if __name__ == "__main__":
    main()