Cosine similarity, dot product, euclidean distance

import numpy as np
from typing import Literal


# ── Базовые реализации ────────────────────────────────────────────

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity между двумя векторами."""
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a, b) / (norm_a * norm_b))


def dot_product(a: np.ndarray, b: np.ndarray) -> float:
    """Скалярное произведение. Для нормализованных = cosine."""
    return float(np.dot(a, b))


def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """Евклидово расстояние. Меньше = похожее."""
    return float(np.linalg.norm(a - b))


def l2_normalize(v: np.ndarray) -> np.ndarray:
    """Нормализовать вектор до единичной длины."""
    norm = np.linalg.norm(v)
    return v / norm if norm > 0 else v


# ── Пример: три метрики на одних данных ──────────────────────────

# Моделируем: два вектора похожего смысла (одно направление),
# но разной длины (разный объём текста)
vec_query = np.array([0.6, 0.8, 0.1, -0.3])          # нормализован: ||q||=1.0
vec_a     = np.array([0.55, 0.75, 0.08, -0.28])       # похожее направление, norm≈0.97
vec_b     = np.array([1.2, 1.6, 0.2, -0.6])           # ТО ЖЕ направление × 2, norm≈2.0
vec_c     = np.array([-0.5, 0.2, 0.8, 0.3])           # другое направление

print("=== Исходные нормы ===")
for name, v in [("query", vec_query), ("A", vec_a), ("B", vec_b), ("C", vec_c)]:
    print(f"  ||{name}|| = {np.linalg.norm(v):.3f}")

print("\n=== Cosine similarity с query ===")
for name, v in [("A", vec_a), ("B", vec_b), ("C", vec_c)]:
    print(f"  cosine(query, {name}) = {cosine_similarity(vec_query, v):.4f}")

print("\n=== Dot product с query ===")
for name, v in [("A", vec_a), ("B", vec_b), ("C", vec_c)]:
    print(f"  dot(query, {name})    = {dot_product(vec_query, v):.4f}")
# ► B получит в 2× больший dot product чем A (хотя направление одинаковое!)

print("\n=== Euclidean distance от query ===")
for name, v in [("A", vec_a), ("B", vec_b), ("C", vec_c)]:
    print(f"  euclidean(query, {name}) = {euclidean_distance(vec_query, v):.4f}")
# ► B далеко от query евклидово, хотя смысл идентичен!

print("\n=== После L2-нормализации: dot == cosine ===")
qn = l2_normalize(vec_query)
for name, v in [("A", vec_a), ("B", vec_b), ("C", vec_c)]:
    vn = l2_normalize(v)
    cos = cosine_similarity(qn, vn)
    dp  = dot_product(qn, vn)
    print(f"  {name}: cosine={cos:.4f}, dot(norm)={dp:.4f}  → {'≡' if abs(cos-dp) < 1e-6 else '≠'}")
# ► cos = dot для нормализованных; B больше не доминирует

import numpy as np
import time


def batch_cosine_search(
    query_vec: np.ndarray,
    doc_matrix: np.ndarray,
    top_k: int = 10,
    pre_normalized: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Найти top-k документов по cosine similarity.

    Args:
        query_vec:       1D вектор запроса, shape (D,)
        doc_matrix:      2D матрица документов, shape (N, D)
        top_k:           Сколько лучших вернуть
        pre_normalized:  True → пропустить нормализацию (быстрее)

    Returns:
        (indices, scores) — индексы и оценки top-k, отсортированные по убыванию
    """
    if not pre_normalized:
        query_vec  = query_vec  / (np.linalg.norm(query_vec) + 1e-8)
        # Нормализуем каждую строку матрицы
        norms      = np.linalg.norm(doc_matrix, axis=1, keepdims=True)
        doc_matrix = doc_matrix / (norms + 1e-8)

    # Матричное умножение: (D,) @ (D, N) = (N,) — все сходства за раз
    scores = doc_matrix @ query_vec         # shape: (N,)

    # argpartition быстрее argsort для нахождения top-k
    if top_k < len(scores):
        top_indices = np.argpartition(scores, -top_k)[-top_k:]
    else:
        top_indices = np.arange(len(scores))

    # Сортируем только top-k
    top_indices = top_indices[np.argsort(scores[top_indices])[::-1]]

    return top_indices, scores[top_indices]


# ── Бенчмарк: наивный цикл vs матричное умножение ──────────────

N = 50_000    # документов
D = 1536      # размерность (text-embedding-3-small)

np.random.seed(42)
docs = np.random.randn(N, D).astype(np.float32)
# Нормализуем заранее (в реальности делается при индексировании)
docs /= np.linalg.norm(docs, axis=1, keepdims=True)

query = np.random.randn(D).astype(np.float32)
query /= np.linalg.norm(query)

# Наивный цикл
t0 = time.time()
naive_scores = [float(np.dot(docs[i], query)) for i in range(N)]
t_naive = time.time() - t0

# Матричное умножение
t0 = time.time()
indices, scores = batch_cosine_search(query, docs, top_k=10, pre_normalized=True)
t_batch = time.time() - t0

print(f"Наивный цикл: {t_naive*1000:.1f} мс")
print(f"Матричное умножение: {t_batch*1000:.1f} мс")
print(f"Ускорение: {t_naive/t_batch:.0f}×")
print(f"\nTop-3 результаты:")
for idx, score in zip(indices[:3], scores[:3]):
    print(f"  doc[{idx:5d}]  cosine = {score:.4f}")

from dataclasses import dataclass
from typing import Optional
import numpy as np


@dataclass
class SearchResult:
    index:    int
    score:    float
    document: str


def search_with_threshold(
    query_vec: np.ndarray,
    doc_matrix: np.ndarray,
    documents: list[str],
    top_k: int = 5,
    min_score: float = 0.70,
) -> list[SearchResult]:
    """
    Поиск с фильтрацией по минимальному порогу cosine.

    min_score: документы с cosine < min_score отбрасываются.
    Если подходящих нет — возвращает пустой список.
    """
    # Нормализуем
    q_norm  = query_vec / np.linalg.norm(query_vec)
    norms   = np.linalg.norm(doc_matrix, axis=1, keepdims=True)
    d_norm  = doc_matrix / (norms + 1e-8)

    scores = d_norm @ q_norm                          # (N,)
    ranked = np.argsort(scores)[::-1][:top_k]         # индексы top-k

    results = []
    for idx in ranked:
        if scores[idx] < min_score:
            break                                     # отсортированы → дальше хуже
        results.append(SearchResult(
            index=int(idx),
            score=float(scores[idx]),
            document=documents[idx],
        ))

    return results


def calibrate_threshold(
    query_vecs: np.ndarray,      # (Q, D) — несколько тестовых запросов
    doc_matrix: np.ndarray,      # (N, D) — корпус
    relevant_ids: list[list[int]], # relevant_ids[q] = индексы релевантных документов
    percentile: float = 10.0,    # нижний percentile среди релевантных пар
) -> float:
    """
    Калибровка минимального порога на тестовых данных.
    Возвращает cosine similarity ниже которого отрезается нерелевантное.
    """
    q_norm = query_vecs / np.linalg.norm(query_vecs, axis=1, keepdims=True)
    d_norm = doc_matrix / np.linalg.norm(doc_matrix, axis=1, keepdims=True)

    relevant_scores = []
    for q_idx, rel_ids in enumerate(relevant_ids):
        qv = q_norm[q_idx]
        for doc_id in rel_ids:
            score = float(d_norm[doc_id] @ qv)
            relevant_scores.append(score)

    threshold = float(np.percentile(relevant_scores, percentile))
    print(f"Релевантных пар: {len(relevant_scores)}")
    print(f"Cosine p10: {threshold:.3f}, p50: {np.median(relevant_scores):.3f}")
    print(f"Рекомендуемый min_score: {threshold:.2f}")
    return threshold


# Пример: порог 0.70 — стандартная отправная точка,
# но всегда калибруйте на своём корпусе!
# Для технической документации: часто 0.75–0.80
# Для новостей / разнообразного текста: 0.60–0.70

"""
Метрики в популярных векторных БД.
Выбор делается ОДИН РАЗ при создании коллекции.
"""

# ── ChromaDB ─────────────────────────────────────────────────────
import chromadb

client = chromadb.Client()

# Доступные метрики в Chroma: "cosine" | "l2" | "ip"
# ip = inner product = dot product
# По умолчанию: "l2" (euclidean²), но для embedding-поиска лучше cosine

cosine_collection = client.create_collection(
    name="docs_cosine",
    metadata={"hnsw:space": "cosine"},   # ← задаём здесь
)

ip_collection = client.create_collection(
    name="docs_ip",
    metadata={"hnsw:space": "ip"},       # для нормализованных = cosine, быстрее
)


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

qdrant = QdrantClient(":memory:")

qdrant.create_collection(
    collection_name="docs",
    vectors_config=VectorParams(
        size=1536,
        distance=Distance.COSINE,      # или DOT / EUCLID / MANHATTAN
    ),
)

# Qdrant нормализует векторы при cosine — хранит unit vectors внутри
# и использует dot product для скорости.
# При Distance.DOT — не нормализует, raw dot product.


# ── FAISS (локально, без сервера) ─────────────────────────────────
import faiss
import numpy as np

D = 1536

# IndexFlatIP = inner product (= cosine для нормализованных)
index_ip = faiss.IndexFlatIP(D)

# IndexFlatL2 = euclidean distance
index_l2 = faiss.IndexFlatL2(D)

# HNSW + IP (approx nearest neighbor, production-grade)
index_hnsw = faiss.IndexHNSWFlat(D, 32)   # 32 = M (graph connectivity)
index_hnsw.metric_type = faiss.METRIC_INNER_PRODUCT

# Пример: индексирование нормализованных векторов
N = 10_000
docs = np.random.randn(N, D).astype(np.float32)
faiss.normalize_L2(docs)                   # нормализуем inplace
index_ip.add(docs)

# Поиск: возвращает (scores, indices)
query = np.random.randn(1, D).astype(np.float32)
faiss.normalize_L2(query)
scores, indices = index_ip.search(query, k=10)
# scores: cosine similarity (т.к. оба нормализованы + IP)