"""
RBSP HNSW -- Hierarchical Navigable Small World Graph for ANN Search.

Pure-Python HNSW with O(log n) search.  Supports cosine and L2 distance.
Binary quantization option for ultra-fast Hamming distance search.

The HNSW graph is a multi-layer proximity graph where:
    - Layer 0 contains ALL nodes (the base layer).
    - Higher layers contain exponentially fewer nodes (geometric decay).
    - Search starts at the topmost layer and greedily descends.
    - At each layer a beam search (width ``ef``) finds nearest neighbours.

Thread-safe via ``threading.Lock``.  Zero external dependencies.
"""

from __future__ import annotations

import heapq
import math
import random
import threading
from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple

# ---------------------------------------------------------------------------
# Distance functions
# ---------------------------------------------------------------------------


def dot_product(a: Sequence[float], b: Sequence[float]) -> float:
    """Compute the dot product of two vectors.

    Args:
        a: First vector.
        b: Second vector.

    Returns:
        Scalar dot product.

    Raises:
        ValueError: If vector lengths differ.
    """
    if len(a) != len(b):
        raise ValueError(
            f"Vector lengths must match: {len(a)} != {len(b)}"
        )
    total = 0.0
    for ai, bi in zip(a, b):
        total += ai * bi
    return total


def normalize_vector(v: Sequence[float]) -> List[float]:
    """Return a unit-length copy of *v*.

    If *v* is the zero vector the original values are returned unchanged.

    Args:
        v: Input vector.

    Returns:
        List of floats with L2 norm == 1.0 (or all zeros).
    """
    norm = math.sqrt(sum(x * x for x in v))
    if norm == 0.0:
        return list(v)
    inv = 1.0 / norm
    return [x * inv for x in v]


def l2_distance(a: Sequence[float], b: Sequence[float]) -> float:
    """Compute the Euclidean (L2) distance between two vectors.

    Args:
        a: First vector.
        b: Second vector.

    Returns:
        Non-negative float distance.

    Raises:
        ValueError: If vector lengths differ.
    """
    if len(a) != len(b):
        raise ValueError(
            f"Vector lengths must match: {len(a)} != {len(b)}"
        )
    total = 0.0
    for ai, bi in zip(a, b):
        d = ai - bi
        total += d * d
    return math.sqrt(total)


def cosine_distance(a: Sequence[float], b: Sequence[float]) -> float:
    """Compute the cosine distance (1 - cosine_similarity) between two vectors.

    Returns a value in [0.0, 2.0] where 0 means identical direction and
    2 means opposite direction.  Returns 1.0 if either vector is zero.

    Args:
        a: First vector.
        b: Second vector.

    Returns:
        Cosine distance in [0.0, 2.0].

    Raises:
        ValueError: If vector lengths differ.
    """
    if len(a) != len(b):
        raise ValueError(
            f"Vector lengths must match: {len(a)} != {len(b)}"
        )
    dot = 0.0
    norm_a = 0.0
    norm_b = 0.0
    for ai, bi in zip(a, b):
        dot += ai * bi
        norm_a += ai * ai
        norm_b += bi * bi
    if norm_a == 0.0 or norm_b == 0.0:
        return 1.0
    sim = dot / (math.sqrt(norm_a) * math.sqrt(norm_b))
    # Clamp to [-1, 1] to handle floating-point rounding
    sim = max(-1.0, min(1.0, sim))
    return 1.0 - sim


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


class HNSWConfig:
    """Configuration for an HNSW index.

    Args:
        M:               Max bi-directional connections per node per layer
                         (default 16).
        ef_construction: Beam width during index construction (default 200).
        ef_search:       Beam width during query (default 50).
        max_level:       Hard cap on graph levels.  When ``None`` (default)
                         the maximum level is computed automatically from
                         ``ln(N) / ln(M)`` during insertion.
        distance:        Distance metric -- ``"cosine"`` or ``"l2"``
                         (default ``"cosine"``).
        seed:            RNG seed for reproducible level assignment.
    """

    __slots__ = (
        "M",
        "ef_construction",
        "ef_search",
        "max_level",
        "distance",
        "seed",
    )

    def __init__(
        self,
        *,
        M: int = 16,
        ef_construction: int = 200,
        ef_search: int = 50,
        max_level: Optional[int] = None,
        distance: str = "cosine",
        seed: int = 42,
    ) -> None:
        if M < 2:
            raise ValueError("M must be >= 2")
        if ef_construction < 1:
            raise ValueError("ef_construction must be >= 1")
        if ef_search < 1:
            raise ValueError("ef_search must be >= 1")
        if distance not in ("cosine", "l2"):
            raise ValueError(f"distance must be 'cosine' or 'l2', got {distance!r}")
        self.M = M
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        self.max_level = max_level
        self.distance = distance
        self.seed = seed

    def __repr__(self) -> str:
        return (
            f"HNSWConfig(M={self.M}, ef_construction={self.ef_construction}, "
            f"ef_search={self.ef_search}, distance={self.distance!r})"
        )


# ---------------------------------------------------------------------------
# Internal node
# ---------------------------------------------------------------------------


class _HNSWNode:
    """Internal graph node holding a vector and per-layer neighbour lists."""

    __slots__ = ("id", "vector", "neighbors", "level")

    def __init__(
        self,
        node_id: str,
        vector: Tuple[float, ...],
        level: int,
    ) -> None:
        self.id = node_id
        self.vector = vector
        self.level = level
        # neighbors[layer] -> list of node IDs
        self.neighbors: List[List[str]] = [[] for _ in range(level + 1)]


# ---------------------------------------------------------------------------
# HNSW Index
# ---------------------------------------------------------------------------


class HNSWIndex:
    """Hierarchical Navigable Small World graph for approximate nearest neighbor search.

    Thread-safe.  All public methods acquire an internal lock.

    Supports:
        - ``insert(id, vector)``          -- add or update a vector
        - ``search(query_vector, k=10)``  -- approximate k-NN query
        - ``remove(id)``                  -- delete a node
        - ``__len__()``                   -- number of indexed nodes
        - ``__contains__(id)``            -- membership test

    Args:
        config: An ``HNSWConfig`` instance (optional, defaults apply).

    Example::

        idx = HNSWIndex()
        idx.insert("a", [0.1, 0.2, 0.3])
        idx.insert("b", [0.4, 0.5, 0.6])
        results = idx.search([0.1, 0.2, 0.3], k=1)
        # results == [("a", 0.0)]
    """

    def __init__(self, config: Optional[HNSWConfig] = None) -> None:
        self._config = config or HNSWConfig()
        self._nodes: Dict[str, _HNSWNode] = {}
        self._entry_point: Optional[str] = None
        self._max_level = 0
        self._lock = threading.Lock()
        self._rng = random.Random(self._config.seed)

        # Select the distance function once
        if self._config.distance == "cosine":
            self._dist = cosine_distance
        else:
            self._dist = l2_distance

        # M_max0 is the max connections for layer 0 (typically 2*M)
        self._M_max0 = self._config.M * 2
        # _m_L controls level probability: 1 / ln(M)
        self._m_L = 1.0 / math.log(self._config.M)

    # -- Public API ---------------------------------------------------------

    def insert(self, node_id: str, vector: Sequence[float]) -> None:
        """Insert a vector into the index.

        If *node_id* already exists the old entry is replaced.

        Args:
            node_id: Unique string identifier for the vector.
            vector:  The feature vector (any sequence of floats).
        """
        vec = tuple(float(x) for x in vector)
        with self._lock:
            if node_id in self._nodes:
                self._remove_node(node_id)
            self._insert_node(node_id, vec)

    def search(
        self,
        query_vector: Sequence[float],
        k: int = 10,
        *,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """Find the *k* approximate nearest neighbors of *query_vector*.

        Args:
            query_vector: The query feature vector.
            k:            Number of neighbours to return (default 10).
            filter_fn:    Optional predicate ``(node_id) -> bool``.  Only
                          nodes for which *filter_fn* returns ``True`` are
                          included in results.  The graph is still traversed
                          through filtered-out nodes to maintain connectivity.

        Returns:
            List of ``(node_id, distance)`` tuples sorted by ascending
            distance.  May return fewer than *k* results if the index
            contains fewer nodes or too few pass the filter.
        """
        qv = tuple(float(x) for x in query_vector)
        with self._lock:
            return self._search(qv, k, filter_fn)

    def remove(self, node_id: str) -> bool:
        """Remove a node from the index.

        Returns:
            ``True`` if the node was found and removed, ``False`` otherwise.
        """
        with self._lock:
            return self._remove_node(node_id)

    def __len__(self) -> int:
        return len(self._nodes)

    def __contains__(self, node_id: str) -> bool:  # type: ignore[override]
        return node_id in self._nodes

    def __repr__(self) -> str:
        return f"HNSWIndex(len={len(self._nodes)}, config={self._config!r})"

    # -- Internal: insertion ------------------------------------------------

    def _select_level(self) -> int:
        """Pick a random level for a new node (exponential distribution).

        Level l is chosen with probability  exp(-l / m_L) * (1 - exp(-1 / m_L)).
        In practice we draw a uniform [0, 1) and take floor(-ln(u) * m_L).
        """
        u = self._rng.random()
        # Avoid log(0)
        if u == 0.0:
            u = 1e-15
        level = int(-math.log(u) * self._m_L)
        if self._config.max_level is not None:
            level = min(level, self._config.max_level)
        return level

    def _insert_node(self, node_id: str, vector: Tuple[float, ...]) -> None:
        """Core insertion logic (caller holds the lock)."""
        level = self._select_level()
        node = _HNSWNode(node_id, vector, level)
        self._nodes[node_id] = node

        # First node -- just set as entry point
        if self._entry_point is None:
            self._entry_point = node_id
            self._max_level = level
            return

        ep = self._entry_point
        ep_node = self._nodes[ep]

        # Phase 1: greedily traverse from the top level down to level+1
        current_ep = ep
        for lc in range(self._max_level, level, -1):
            current_ep = self._search_layer_greedy(vector, current_ep, lc)

        # Phase 2: at each layer from min(level, _max_level) down to 0
        #          do a beam search and connect the new node
        ef = self._config.ef_construction
        candidates_ep: Set[str] = {current_ep}
        for lc in range(min(level, self._max_level), -1, -1):
            candidates = self._search_layer(vector, candidates_ep, ef, lc)
            m_max = self._M_max0 if lc == 0 else self._config.M
            neighbors = self._select_neighbors(vector, candidates, m_max, lc)

            # Connect new node -> selected neighbours
            node.neighbors[lc] = list(neighbors)

            # Connect neighbours -> new node (bi-directional)
            for nbr_id in neighbors:
                nbr_node = self._nodes[nbr_id]
                nbr_node.neighbors[lc].append(node_id)
                # Prune if too many connections
                if len(nbr_node.neighbors[lc]) > m_max:
                    nbr_node.neighbors[lc] = self._select_neighbors(
                        nbr_node.vector,
                        set(nbr_node.neighbors[lc]),
                        m_max,
                        lc,
                    )

            # For next layer, use the found candidates as entry points
            candidates_ep = set(neighbors) if neighbors else candidates_ep

        # Update entry point if the new node reaches a higher level
        if level > self._max_level:
            self._entry_point = node_id
            self._max_level = level

    # -- Internal: search ---------------------------------------------------

    def _search(
        self,
        query: Tuple[float, ...],
        k: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, float]]:
        """Core search logic (caller holds the lock).

        When *filter_fn* is provided the beam width is tripled and the
        filter is applied **inside** ``_search_layer`` so that only
        matching nodes consume results-heap slots.  Non-matching nodes
        are still traversed via the candidates heap to preserve graph
        connectivity.
        """
        if self._entry_point is None:
            return []

        ep = self._entry_point

        # Phase 1: greedy descent from top to layer 1
        for lc in range(self._max_level, 0, -1):
            ep = self._search_layer_greedy(query, ep, lc)

        # Phase 2: beam search at layer 0
        # Widen ef when filtering so we explore more of the graph
        ef = max(self._config.ef_search, k)
        if filter_fn is not None:
            ef = ef * 3
        candidates = self._search_layer(query, {ep}, ef, 0, filter_fn)

        # Return top-k sorted by distance ascending
        scored: List[Tuple[float, str]] = []
        for cid in candidates:
            d = self._dist(query, self._nodes[cid].vector)
            scored.append((d, cid))

        top = heapq.nsmallest(k, scored, key=lambda x: x[0])
        return [(nid, dist) for dist, nid in top]

    def _search_layer_greedy(
        self,
        query: Tuple[float, ...],
        entry_point: str,
        layer: int,
    ) -> str:
        """Greedy traversal: find the single nearest node at *layer*.

        Used during the top-down descent in both insert and search.
        """
        current = entry_point
        current_dist = self._dist(query, self._nodes[current].vector)

        changed = True
        while changed:
            changed = False
            node = self._nodes[current]
            if layer < len(node.neighbors):
                for nbr_id in node.neighbors[layer]:
                    if nbr_id not in self._nodes:
                        continue
                    d = self._dist(query, self._nodes[nbr_id].vector)
                    if d < current_dist:
                        current = nbr_id
                        current_dist = d
                        changed = True

        return current

    def _search_layer(
        self,
        query: Tuple[float, ...],
        entry_points: Set[str],
        ef: int,
        layer: int,
        filter_fn: Optional[Callable[[str], bool]] = None,
    ) -> Set[str]:
        """Beam search at a single layer.

        Returns a set of up to *ef* nearest node IDs at *layer*.

        Uses two heaps:
            - ``candidates`` (min-heap): nodes still to explore, ordered
              by distance ascending.  **All** nodes are added here to
              preserve graph connectivity during filtered search.
            - ``results`` (max-heap): best nodes found so far, ordered by
              *negative* distance so the worst result is at the top.
              When *filter_fn* is given, only nodes that pass the filter
              are added to results.

        Args:
            query:        Query vector.
            entry_points: Starting node IDs.
            ef:           Beam width.
            layer:        Graph layer to search.
            filter_fn:    Optional predicate ``(node_id) -> bool``.
        """
        visited: Set[str] = set(entry_points)

        # candidates: (distance, node_id) -- min-heap
        candidates: List[Tuple[float, str]] = []
        # results: (-distance, node_id) -- max-heap (worst at top)
        results: List[Tuple[float, str]] = []

        for ep in entry_points:
            if ep not in self._nodes:
                continue
            d = self._dist(query, self._nodes[ep].vector)
            heapq.heappush(candidates, (d, ep))
            if filter_fn is None or filter_fn(ep):
                heapq.heappush(results, (-d, ep))

        # If results is empty after entry points (all filtered), we still
        # need a distance bound; use infinity so exploration continues.
        while candidates:
            c_dist, c_id = heapq.heappop(candidates)
            # Furthest result distance (or inf if no results yet)
            if results:
                f_dist = -results[0][0]
            else:
                f_dist = float("inf")
            if c_dist > f_dist and len(results) >= ef:
                break

            c_node = self._nodes.get(c_id)
            if c_node is None:
                continue
            if layer >= len(c_node.neighbors):
                continue

            for nbr_id in c_node.neighbors[layer]:
                if nbr_id in visited:
                    continue
                visited.add(nbr_id)
                if nbr_id not in self._nodes:
                    continue

                d = self._dist(query, self._nodes[nbr_id].vector)
                if results:
                    f_dist = -results[0][0]
                else:
                    f_dist = float("inf")

                if d < f_dist or len(results) < ef:
                    # Always add to candidates for navigation
                    heapq.heappush(candidates, (d, nbr_id))
                    # Only add to results if filter passes
                    if filter_fn is None or filter_fn(nbr_id):
                        heapq.heappush(results, (-d, nbr_id))
                        if len(results) > ef:
                            heapq.heappop(results)

        return {nid for _, nid in results}

    # -- Internal: neighbour selection --------------------------------------

    def _select_neighbors(
        self,
        target_vector: Tuple[float, ...],
        candidates: Set[str],
        M: int,
        layer: int,
    ) -> List[str]:
        """Select up to *M* neighbours from *candidates* for a node at *target_vector*.

        Uses the simple selection strategy: pick the *M* closest candidates
        by distance to *target_vector*.
        """
        if not candidates:
            return []

        scored: List[Tuple[float, str]] = []
        for cid in candidates:
            node = self._nodes.get(cid)
            if node is None:
                continue
            d = self._dist(target_vector, node.vector)
            scored.append((d, cid))

        selected = heapq.nsmallest(M, scored, key=lambda x: x[0])
        return [nid for _, nid in selected]

    # -- Internal: removal --------------------------------------------------

    def _remove_node(self, node_id: str) -> bool:
        """Remove a node and all its connections (caller holds the lock)."""
        node = self._nodes.pop(node_id, None)
        if node is None:
            return False

        # Remove all inbound references from neighbours
        for lc in range(node.level + 1):
            for nbr_id in node.neighbors[lc]:
                nbr = self._nodes.get(nbr_id)
                if nbr is not None and lc < len(nbr.neighbors):
                    try:
                        nbr.neighbors[lc].remove(node_id)
                    except ValueError:
                        pass

        # If the removed node was the entry point, pick a replacement
        if self._entry_point == node_id:
            if not self._nodes:
                self._entry_point = None
                self._max_level = 0
            else:
                # Find the node with the highest level
                best_id: Optional[str] = None
                best_level = -1
                for nid, n in self._nodes.items():
                    if n.level > best_level:
                        best_level = n.level
                        best_id = nid
                self._entry_point = best_id
                self._max_level = best_level

        return True
