"""Small vector index for nearest-neighbor exploration."""

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True, slots=True)
class Neighbor:
    """A nearest-neighbor result from a vector search."""

    item_id: str
    score: float


class VectorIndex:
    """In-memory cosine-similarity vector index."""

    def __init__(self) -> None:
        """Create an empty vector index."""

        self._vectors: dict[str, tuple[float, ...]] = {}

    def add(self, item_id: str, vector: tuple[float, ...]) -> None:
        """Add or replace an item vector."""

        if not vector:
            raise ValueError("Vector must contain at least one value.")
        if not all(math.isfinite(value) for value in vector):
            raise ValueError("Vector values must be finite.")
        self._vectors[item_id] = vector

    def search(
        self, query_vector: tuple[float, ...], *, limit: int = 5, exclude_id: str | None = None
    ) -> list[Neighbor]:
        """Return the closest vectors by cosine similarity."""

        if limit < 1:
            raise ValueError("Limit must be at least 1.")
        neighbors = [
            Neighbor(item_id=item_id, score=_cosine_similarity(query_vector, vector))
            for item_id, vector in self._vectors.items()
            if item_id != exclude_id
        ]
        return sorted(neighbors, key=lambda neighbor: neighbor.score, reverse=True)[:limit]

    def save_tsv(self, path: Path) -> None:
        """Persist the index as a simple tab-separated text file."""

        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as index_file:
            for item_id, vector in sorted(self._vectors.items()):
                values = "\t".join(str(value) for value in vector)
                index_file.write(f"{item_id}\t{values}\n")

    @classmethod
    def load_tsv(cls, path: Path) -> VectorIndex:
        """Load an index produced by :meth:`save_tsv`."""

        index = cls()
        with path.open("r", encoding="utf-8") as index_file:
            for line_number, line in enumerate(index_file, start=1):
                fields = line.rstrip("\n").split("\t")
                if len(fields) < 2:
                    msg = f"Line {line_number} must contain an id and vector values."
                    raise ValueError(msg)
                index.add(fields[0], tuple(float(value) for value in fields[1:]))
        return index

    def __len__(self) -> int:
        """Return the number of indexed vectors."""

        return len(self._vectors)


def _cosine_similarity(left: tuple[float, ...], right: tuple[float, ...]) -> float:
    if len(left) != len(right):
        raise ValueError("Vectors must have the same dimensions.")
    left_norm = math.sqrt(sum(value * value for value in left))
    right_norm = math.sqrt(sum(value * value for value in right))
    if left_norm == 0.0 or right_norm == 0.0:
        return 0.0
    dot_product = sum(
        left_value * right_value for left_value, right_value in zip(left, right, strict=True)
    )
    return dot_product / (left_norm * right_norm)


def cosine_similarity(left: tuple[float, ...], right: tuple[float, ...]) -> float:
    """Public wrapper for cosine similarity (useful in notebooks/examples)."""

    return _cosine_similarity(left, right)