import numpy as np

a = np.array([0.7, 0.8, 0.9])
b = np.array([0.6, 0.7, 0.8])

# ① 欧氏距离（Euclidean）—— 直线距离
#    越小越相似，受向量长度影响大
euclidean = np.linalg.norm(a - b)
print(f"欧氏距离: {euclidean:.4f}")  # ~0.173

# ② 余弦相似度（Cosine）—— 夹角
#    越大越相似（范围 -1 到 1），不受向量长度影响
cosine = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
print(f"余弦相似度: {cosine:.4f}")  # ~0.997

# ③ 内积（Dot Product）—— 投影长度
#    越大越相似，某种程度上等价于余弦相似度
dot = np.dot(a, b)
print(f"内积: {dot:.4f}")  # ~1.70

# IVF 工作流程（简化实现）
from sklearn.cluster import KMeans
import numpy as np

class SimpleIVF:
    def __init__(self, n_clusters=1000, n_probe=10):
        self.n_clusters = n_clusters
        self.n_probe = n_probe  # 查询时要探索的聚类数
        self.kmeans = KMeans(n_clusters=n_clusters)
        self.vectors = None
        self.cluster_assignments = None
    
    def build(self, vectors):
        """构建索引"""
        self.vectors = vectors
        self.cluster_assignments = self.kmeans.fit_predict(vectors)
        # 每个聚类存一个向量列表
        self.clusters = {i: [] for i in range(self.n_clusters)}
        for idx, cluster_id in enumerate(self.cluster_assignments):
            self.clusters[cluster_id].append(idx)
    
    def search(self, query_vec, k=10):
        """搜索 Top-K"""
        # Step 1: 找到最近的 n_probe 个聚类中心
        centers = self.kmeans.cluster_centers_
        distances = np.linalg.norm(centers - query_vec, axis=1)
        nearest_clusters = np.argsort(distances)[:self.n_probe]
        
        # Step 2: 在这几个聚类中暴力检索
        candidates = []
        for cid in nearest_clusters:
            for idx in self.clusters[cid]:
                dist = np.linalg.norm(self.vectors[idx] - query_vec)
                candidates.append((idx, dist))
        
        # Step 3: 排序取 Top-K
        candidates.sort(key=lambda x: x[1])
        return candidates[:k]

# 使用示例
ivf = SimpleIVF(n_clusters=1000, n_probe=10)
ivf.build(my_vectors)          # 100 万向量
results = ivf.search(query, k=10)  # 只检索约 10,000 个向量（原来的 1/100）

# HNSW 核心参数
HNSW_CONFIG = {
    "M": 16,            # 每个节点最多有多少个"朋友"（边数）
                        # 建议范围：4-64，越大精度越高但内存越大
    "ef_construction": 200,  # 构建时的搜索宽度
                              # 建议范围：100-500，越大构建越慢但图越好
    "ef_search": 50,    # 查询时的搜索宽度
                        # 建议范围：50-200，越大精度越高但越慢
}

# M 的影响：
# M=8  → 每个节点 8 个朋友  → 内存小，精度 ~90%
# M=16 → 每个节点 16 个朋友 → 内存中，精度 ~95%（最常用）
# M=64 → 每个节点 64 个朋友 → 内存大，精度 ~99%

# ef_construction 和 ef_search 的区别：
# ef_construction：构建时用的，调大 = 更好的图 = 更慢的构建
# ef_search：查询时用的，调大 = 更好的召回 = 更慢的查询

# PQ 压缩示例
from sklearn.cluster import KMeans
import numpy as np

class SimplePQ:
    def __init__(self, n_subvectors=64, n_centroids=256):
        self.n_subvectors = n_subvectors  # 分成几段
        self.n_centroids = n_centroids    # 每段几个聚类中心
        self.codebooks = []               # 64 个"词典"
    
    def train(self, vectors):
        """训练 64 个 K-Means 模型"""
        dim = vectors.shape[1]
        sub_dim = dim // self.n_subvectors
        
        for i in range(self.n_subvectors):
            start = i * sub_dim
            end = (i + 1) * sub_dim
            sub_vectors = vectors[:, start:end]
            
            kmeans = KMeans(n_clusters=self.n_centroids)
            kmeans.fit(sub_vectors)
            self.codebooks.append(kmeans)
    
    def compress(self, vector):
        """把一个 1024 维向量压缩成 64 个编码"""
        dim = len(vector)
        sub_dim = dim // self.n_subvectors
        codes = []
        
        for i in range(self.n_subvectors):
            start = i * sub_dim
            end = (i + 1) * sub_dim
            sub = vector[start:end].reshape(1, -1)
            code = self.codebooks[i].predict(sub)[0]
            codes.append(code)
        
        return np.array(codes, dtype=np.uint8)  # 64 字节

# Milvus 使用示例
from pymilvus import MilvusClient, DataType

client = MilvusClient("http://localhost:19530")

# 创建集合（需要定义 Schema）
schema = client.create_schema(
    auto_id=False,
    enable_dynamic_field=True
)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("text", DataType.VARCHAR, max_length=65535)
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=1024)

client.create_collection(
    collection_name="my_docs",
    schema=schema,
    index_params={
        "field_name": "vector",
        "index_type": "HNSW",
        "metric_type": "COSINE",
        "params": {"M": 16, "efConstruction": 200}
    }
)

# 插入
client.insert("my_docs", [
    {"id": 1, "text": "人工智能正在改变世界...", "vector": embedding_1},
    {"id": 2, "text": "机器学习是 AI 的子领域...", "vector": embedding_2},
])

# 搜索
results = client.search(
    collection_name="my_docs",
    data=[query_embedding],
    limit=10,
    output_fields=["text"]  # 返回文本字段
)

from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, Range

client = QdrantClient("localhost", port=6333)

# 精确到令人发指的过滤搜索
results = client.search(
    collection_name="products",
    query_vector=query_emb,
    limit=10,
    query_filter=Filter(
        must=[
            FieldCondition(
                key="price",
                range=Range(gte=100, lte=500)  # 价格 100-500
            ),
            FieldCondition(
                key="category",
                match={"value": "electronics"}  # 分类：电子产品
            ),
        ],
        must_not=[
            FieldCondition(
                key="brand",
                match={"value": "山寨牌"}  # 排除山寨品牌
            )
        ]
    )
)

import lancedb
import numpy as np

# 零配置，直接开始！
db = lancedb.connect("./my_lancedb")  # 一个目录 = 一个数据库

# 创建表
table = db.create_table(
    "documents",
    data=[{
        "vector": np.random.randn(1024).tolist(),
        "text": "LanceDB 太方便了",
        "id": 1
    }]
)

# 创建索引（支持 IVF_PQ, IVF_HNSW_PQ, DiskANN）
table.create_index(
    metric="cosine",
    num_partitions=256,
    num_sub_vectors=64  # PQ 压缩参数
)

# 搜索
results = table.search(query_vec).limit(10).to_pandas()

import chromadb

# 最简洁的 API（之一）
client = chromadb.PersistentClient(path="./chroma_data")
collection = client.get_or_create_collection("my_docs")

# 插入（自动生成 ID）
collection.add(
    documents=["这是一段文本", "另一段文本"],
    embeddings=[[0.1, 0.2, ...], [0.3, 0.4, ...]],
    metadatas=[{"source": "a.pdf"}, {"source": "b.pdf"}],
    ids=["doc1", "doc2"]
)

# 查询（甚至可以直接用文本，内置了默认的 Embedding 模型！）
results = collection.query(
    query_texts=["什么是向量数据库？"],
    n_results=10
)
# ChromaDB 有一个内置的 all-MiniLM-L6-v2 Embedding（不推荐生产用）

import weaviate

client = weaviate.Client("http://localhost:8080")

# 什么都不用管，直接存入文本
client.data_object.create(
    data_object={"content": "这是一篇关于 AI 的文章"},
    class_name="Document",
    # Weaviate 自动调 Embedding 服务把文本转成向量！
)

# GraphQL 查询（Weaviate 用 GraphQL 而不是 REST）
response = client.query.get(
    "Document", ["content"]
).with_near_text({
    "concepts": ["人工智能"]  # 自动把查询文本转成向量
}).with_limit(10).do()

# 1. 批量生成 Embedding（而非逐条）
# ❌ 坏做法——每次 encode 一条
for doc in documents:
    embedding = model.encode(doc)  # GPU 利用率极低
    db.insert(embedding)

# ✅ 好做法——批量 encode
batch_size = 256
for i in range(0, len(documents), batch_size):
    batch = documents[i:i+batch_size]
    embeddings = model.encode(batch, show_progress_bar=False)
    db.insert_batch(embeddings, batch)

# 2. 向量归一化（让内积 = 余弦相似度）
# 大多数 Embedding 模型输出未归一化，需要：
import torch.nn.functional as F
normalized_emb = F.normalize(torch.tensor(emb), p=2, dim=-1).numpy()

# 3. 使用 FP16 半精度（精度损失可忽略，内存减半）
model = SentenceTransformer("BAAI/bge-m3", device="cuda")
model.half()  # float32 → float16

# 多向量表示例
item = {
    "id": "product_12345",
    "title": "超级好用的蓝牙耳机",
    "text_embedding": [0.1, 0.2, ...],      # 文本语义向量
    "image_embedding": [-0.3, 0.5, ...],    # 图片视觉向量
    "audio_embedding": [0.8, -0.1, ...],    # 音频向量（如果有）
    "late_interaction_embedding": [         # ColBERT 风格的多向量
        [0.1, 0.2, ...],  # token 1
        [0.3, 0.4, ...],  # token 2
        ...
    ]
}
# 一条记录有多个向量，检索时可以同时考虑

"""
向量数据库 Benchmark
对比 ChromaDB / Qdrant / LanceDB 的插入和查询性能
"""
import time
import numpy as np
from typing import List, Tuple

# ─── 生成 10 万个 1024 维向量 ───
N = 100_000
DIM = 1024

np.random.seed(42)
train_vectors = np.random.randn(N, DIM).astype(np.float32)
# 归一化
train_vectors = train_vectors / np.linalg.norm(train_vectors, axis=1, keepdims=True)
query_vectors = np.random.randn(100, DIM).astype(np.float32)
query_vectors = query_vectors / np.linalg.norm(query_vectors, axis=1, keepdims=True)

# ─── 1. ChromaDB ───
import chromadb

chroma_client = chromadb.PersistentClient(path="./bench_chroma")
chroma_collection = chroma_client.get_or_create_collection(
    "bench", metadata={"hnsw:space": "cosine"}
)

# 插入
t0 = time.time()
batch_size = 5000
for i in range(0, N, batch_size):
    batch = train_vectors[i:i+batch_size]
    chroma_collection.add(
        embeddings=batch.tolist(),
        ids=[f"vec_{j}" for j in range(i, min(i+batch_size, N))]
    )
chroma_insert_time = time.time() - t0

# 查询
t0 = time.time()
for qv in query_vectors:
    chroma_collection.query(query_embeddings=[qv.tolist()], n_results=10)
chroma_query_time = time.time() - t0

# ─── 2. LanceDB ───
import lancedb

lance_db = lancedb.connect("./bench_lancedb")
# 准备数据
data = [{"vector": train_vectors[i].tolist(), "id": i} for i in range(N)]
t0 = time.time()
lance_table = lance_db.create_table("bench", data=data, mode="overwrite")
lance_insert_time = time.time() - t0

# 创建索引
t0 = time.time()
lance_table.create_index(
    metric="cosine",
    num_partitions=256,
    num_sub_vectors=64
)
lance_index_time = time.time() - t0

# 查询
t0 = time.time()
for qv in query_vectors:
    lance_table.search(qv.tolist()).limit(10).to_list()
lance_query_time = time.time() - t0

# ─── 3. Qdrant ───
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

qdrant_client = QdrantClient("localhost", port=6333)

# 删除旧集合
try:
    qdrant_client.delete_collection("bench")
except:
    pass

qdrant_client.create_collection(
    "bench",
    vectors_config=VectorParams(size=DIM, distance=Distance.COSINE)
)

# 插入
t0 = time.time()
points = [
    PointStruct(
        id=i,
        vector=train_vectors[i].tolist()
    )
    for i in range(N)
]
qdrant_client.upsert("bench", points)
qdrant_insert_time = time.time() - t0

# 创建索引
t0 = time.time()
qdrant_client.create_payload_index(
    "bench", field_name="id", field_schema="integer"
)
qdrant_index_time = time.time() - t0

# 查询
t0 = time.time()
for qv in query_vectors:
    qdrant_client.search("bench", query_vector=qv.tolist(), limit=10)
qdrant_query_time = time.time() - t0

# ─── 结果汇总 ───
print(f"""
╔══════════════════════════════════════════╗
║    向量数据库 Benchmark: {N:,} × {DIM}维      ║
╠══════════════════╦══════════╦═══════════╣
║      数据库       ║ 插入耗时  ║  查询耗时   ║
╠══════════════════╬══════════╬═══════════╣
║ ChromaDB         ║ {chroma_insert_time:6.1f}s ║ {chroma_query_time:7.2f}s  ║
║ LanceDB          ║ {lance_insert_time:6.1f}s ║ {lance_query_time:7.2f}s  ║
║ Qdrant           ║ {qdrant_insert_time:6.1f}s ║ {qdrant_query_time:7.2f}s  ║
╚══════════════════╩══════════╩═══════════╝

注：LanceDB 索引构建耗时: {lance_index_time:.1f}s
    Qdrant 索引构建耗时: {qdrant_index_time:.1f}s
    (ChromaDB 索引为自动构建，无单独计时)

预期典型结果（10万×1024维）：
  · 插入：Qdrant ~1s < LanceDB ~2s < ChromaDB ~5s
  · 查询：Qdrant ~0.5s < LanceDB ~1s < ChromaDB ~2s
""")