"""HNSW memory integration for RL agent."""

import hnswlib
import numpy as np
from typing import List, Tuple, Optional


class HNSWMemory:
    """HNSW-based memory for experience retrieval.

    Uses Hierarchical Navigable Small World graphs for efficient
    nearest neighbor search in experience space.
    """

    def __init__(
        self,
        dim: int = 256,
        max_elements: int = 100000,
        M: int = 16,
        ef_construction: int = 200,
        ef_search: int = 50,
    ):
        """Initialize HNSW memory.

        Args:
            dim: State dimension
            max_elements: Maximum number of elements
            M: HNSW M parameter (connections per layer)
            ef_construction: HNSW ef_construction parameter
            ef_search: HNSW ef_search parameter
        """
        self.dim = dim
        self.max_elements = max_elements

        # Initialize HNSW index
        self.index = hnswlib.Index(space='l2', dim=dim)
        self.index.init_index(
            max_elements=max_elements,
            M=M,
            ef_construction=ef_construction,
        )
        self.index.set_ef(ef_search)

        # Storage for experiences
        self.experiences = []
        self.current_size = 0

    def add(self, state: np.ndarray, experience: dict) -> None:
        """Add experience to memory.

        Args:
            state: State vector [dim]
            experience: Experience dictionary
        """
        if self.current_size < self.max_elements:
            self.index.add_items(state.reshape(1, -1), [self.current_size])
            self.experiences.append(experience)
            self.current_size += 1

    def query(self, state: np.ndarray, k: int = 5) -> List[dict]:
        """Query k nearest experiences.

        Args:
            state: Query state vector [dim]
            k: Number of neighbors

        Returns:
            List of k nearest experiences
        """
        if self.current_size == 0:
            return []

        k = min(k, self.current_size)
        labels, distances = self.index.knn_query(state.reshape(1, -1), k=k)

        return [self.experiences[idx] for idx in labels[0]]

    def clear(self) -> None:
        """Clear memory."""
        self.index = hnswlib.Index(space='l2', dim=self.dim)
        self.index.init_index(max_elements=self.max_elements)
        self.experiences = []
        self.current_size = 0

    def __len__(self) -> int:
        """Get memory size."""
        return self.current_size
