"""
ManiscopeEngine: Multi-scale Information Retrieval with Geodesic Reranking

This implements a two-stage retrieval system:
1. Telescope (Coarse): Broad retrieval using cosine similarity
2. Microscope (Fine): Geodesic reranking on k-NN manifold graph

Available Versions:
- ManiscopeEngine (v0): Baseline implementation
- ManiscopeEngine_v1: GPU auto-detection + graph caching (3× speedup)
- ManiscopeEngine_v2: Full optimization with FAISS + scipy (5× speedup)
- ManiscopeEngine_v3: Persistent cache + query cache (20-235× speedup)
- ManiscopeEngine_v2o: **RECOMMENDED** - Ultimate optimization (20-235× speedup)

ManiscopeEngine_v2o combines ALL optimizations:
- GPU auto-detection (v1)
- FAISS + scipy sparse matrices (v2)
- Persistent disk cache + query LRU cache (v3)

Real-World Performance (v2o):
- MS MARCO: 132ms → 0.58ms (229× speedup)
- TREC-COVID: 85ms → 0.38ms (226× speedup)
- SciFact: 92ms → 0.39ms (235× speedup)
- 100% accuracy preservation (MRR=1.0)

Key improvements over baseline:
- Better anchor node selection (from candidate set)
- Hybrid scoring combining cosine + geodesic
- Robust handling of disconnected components
- Configurable parameters for k and retrieval depth
"""

import numpy as np
import networkx as nx
from sklearn.neighbors import NearestNeighbors
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from typing import List, Tuple, Dict, Optional, Any
import hashlib
import pickle
import os
from pathlib import Path
from collections import OrderedDict
from heapq import heappush, heappop
import joblib


class ManiscopeEngine:
    """
    Multi-scale retrieval engine combining cosine similarity (telescope)
    and geodesic distance on k-NN manifold (microscope).
    """

    def __init__(self, model_name: str = 'all-MiniLM-L6-v2', k: int = 5,
                 alpha: float = 0.5, verbose: bool = False, device: str = 'cpu',
                 local_files_only: bool = True):
        """
        Initialize the ManiscopeEngine.

        Args:
            model_name: Sentence transformer model name
            k: Number of nearest neighbors for manifold graph construction
            alpha: Weight for hybrid scoring (0=pure geodesic, 1=pure cosine)
            verbose: Print debug information
            device: Device for computation ('cpu' or 'cuda')
            local_files_only: If True, use only cached models without checking HuggingFace
                             (Recommended for experiments to avoid network issues)
        """
        # Load model with local_files_only to avoid checking HuggingFace on every startup
        # This prevents timeout errors when HuggingFace is slow/unavailable
        # and makes experiments more reliable with consistent cached models
        self.model = SentenceTransformer(model_name, device=device,
                                        local_files_only=local_files_only)
        self.k = k
        self.alpha = alpha
        self.verbose = verbose
        self.documents = []
        self.embeddings = None
        self.G = None

    def fit(self, docs: List[str]):
        """
        Build the manifold graph from document corpus.

        Args:
            docs: List of document texts
        """
        if self.verbose:
            print(f"Encoding {len(docs)} documents...")

        self.documents = docs
        self.embeddings = self.model.encode(docs, show_progress_bar=self.verbose)

        if self.verbose:
            print(f"Building k-NN manifold graph (k={self.k})...")

        # Build the latent k-NN manifold graph
        # Use cosine metric to respect embedding geometry
        nbrs = NearestNeighbors(n_neighbors=min(self.k + 1, len(docs)),
                               metric='cosine').fit(self.embeddings)
        distances, indices = nbrs.kneighbors(self.embeddings)

        self.G = nx.Graph()
        for i in range(len(docs)):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    # Add edge with distance as weight
                    self.G.add_edge(i, neighbor_idx, weight=dist)

        if self.verbose:
            print(f"Graph built: {self.G.number_of_nodes()} nodes, "
                  f"{self.G.number_of_edges()} edges")
            print(f"Connected components: "
                  f"{nx.number_connected_components(self.G)}")

    def search_baseline(self, query: str, top_n: int = 5) -> List[Tuple[str, float, int]]:
        """
        Baseline: Telescope Only (Pure Cosine Similarity)

        Args:
            query: Query text
            top_n: Number of results to return

        Returns:
            List of (document, score, doc_index) tuples
        """
        q_emb = self.model.encode([query])
        sims = cosine_similarity(q_emb, self.embeddings)[0]
        idx = np.argsort(sims)[::-1][:top_n]
        return [(self.documents[i], float(sims[i]), int(i)) for i in idx]

    def search_maniscope(self, query: str, top_n: int = 5,
                         coarse_multiplier: int = 3) -> List[Tuple[str, float, int]]:
        """
        Maniscope: Telescope + Microscope (Geodesic Reranking)

        Args:
            query: Query text
            top_n: Number of final results to return
            coarse_multiplier: Retrieve top_n * coarse_multiplier for reranking

        Returns:
            List of (document, score, doc_index) tuples
        """
        # Phase 1: TELESCOPE - Broad retrieval with cosine similarity
        q_emb = self.model.encode([query])
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        # Retrieve broader set for reranking
        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        if self.verbose:
            print(f"\nTelescope: Retrieved top {coarse_size} candidates")

        # Phase 2: MICROSCOPE - Geodesic reranking on manifold
        # Select anchor node from coarse set (closest to query)
        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]  # Most similar in coarse set

        if self.verbose:
            print(f"Microscope: Anchor node = {anchor_node} "
                  f"(cosine sim = {cosine_sims[anchor_node]:.3f})")

        results = []
        for i in coarse_idx:
            cosine_score = cosine_sims[i]

            # Calculate geodesic distance on manifold
            try:
                geo_dist = nx.shortest_path_length(
                    self.G, source=anchor_node, target=i, weight='weight'
                )
                # Convert distance to similarity score
                geo_score = 1.0 / (1.0 + geo_dist)
            except nx.NetworkXNoPath:
                # Disconnected component - use only cosine
                geo_score = 0.0
                if self.verbose:
                    print(f"  Node {i} disconnected from anchor")

            # Hybrid scoring: combine cosine and geodesic
            # alpha=1: pure cosine, alpha=0: pure geodesic
            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            results.append((self.documents[i], final_score, i, cosine_score, geo_score))

        # Sort by final hybrid score
        results.sort(key=lambda x: x[1], reverse=True)

        # Return top_n with simplified format
        return [(doc, score, idx) for doc, score, idx, _, _ in results[:top_n]]

    def search_maniscope_detailed(self, query: str, top_n: int = 5,
                                  coarse_multiplier: int = 3) -> List[dict]:
        """
        Maniscope search with detailed scoring information for analysis.

        Returns:
            List of dicts with document, scores, and metadata
        """
        q_emb = self.model.encode([query])
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        results = []
        for i in coarse_idx:
            cosine_score = float(cosine_sims[i])

            try:
                geo_dist = nx.shortest_path_length(
                    self.G, source=anchor_node, target=i, weight='weight'
                )
                geo_score = float(1.0 / (1.0 + geo_dist))
                connected = True
            except nx.NetworkXNoPath:
                geo_dist = None
                geo_score = 0.0
                connected = False

            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            results.append({
                'doc_id': int(i),
                'document': self.documents[i],
                'final_score': float(final_score),
                'cosine_score': cosine_score,
                'geo_score': geo_score,
                'geo_distance': geo_dist,
                'connected': connected
            })

        results.sort(key=lambda x: x['final_score'], reverse=True)
        return results[:top_n]

    def compare_methods(self, query: str, top_n: int = 5) -> dict:
        """
        Compare baseline and maniscope methods side-by-side.

        Returns:
            Dict with baseline_results, maniscope_results, and analysis
        """
        baseline = self.search_baseline(query, top_n)
        maniscope = self.search_maniscope_detailed(query, top_n)

        # Check if rankings differ
        baseline_ids = [doc_id for _, _, doc_id in baseline]
        maniscope_ids = [result['doc_id'] for result in maniscope]

        ranking_changed = baseline_ids != maniscope_ids

        return {
            'query': query,
            'baseline_results': [
                {'doc_id': doc_id, 'document': doc[:100], 'score': score}
                for doc, score, doc_id in baseline
            ],
            'maniscope_results': [
                {
                    'doc_id': r['doc_id'],
                    'document': r['document'][:100],
                    'final_score': r['final_score'],
                    'cosine_score': r['cosine_score'],
                    'geo_score': r['geo_score']
                }
                for r in maniscope
            ],
            'ranking_changed': ranking_changed,
            'baseline_top1': baseline_ids[0],
            'maniscope_top1': maniscope_ids[0]
        }


class ManiscopeEngine_v3(ManiscopeEngine):
    """
    v3 Optimization: Persistent Caching + Query Cache + Heap (CPU-friendly)

    renamed from __v1

    Target: Variable speedup depending on cache hit rate

    Key improvements over v0 baseline:
    - Persistent disk cache for embeddings with joblib (survives restarts!)
    - Query embedding LRU cache (100 queries in memory)
    - Batch geodesic computation with NetworkX
    - Heap-based early termination for top-n selection
    - Graph adjacency list pre-computation

    Performance: Highly variable
    - First run (cold cache): ~115ms (same as v0)
    - Cache hit (warm): ~30-50ms (2-3x faster)
    - Repeated queries: ~10-20ms (5-10x faster with query cache)

    Best for: Repeated experiments, grid search, CPU-only machines
    Accuracy: MRR=1.0000 (same as v0, no regression)
    """

    def __init__(
        self,
        model_name: str = 'all-MiniLM-L6-v2',
        k: int = 5,
        alpha: float = 0.5,
        verbose: bool = False,
        device: str = 'cpu',
        local_files_only: bool = True,
        cache_dir: Optional[str] = None,
        use_cache: bool = True,
        query_cache_size: int = 100  # New: LRU cache size for queries
    ):
        """
        Initialize the optimized ManiscopeEngine.

        Args:
            model_name: Sentence transformer model name
            k: Number of nearest neighbors for manifold graph construction
            alpha: Weight for hybrid scoring (0=pure geodesic, 1=pure cosine)
            verbose: Print debug information
            device: Device for computation ('cpu' or 'cuda')
            local_files_only: Use only cached models without checking HuggingFace
            cache_dir: Directory for embedding cache (defaults to ~/projects/embedding_cache/maniscope)
            use_cache: Enable disk-based caching of embeddings
            query_cache_size: Maximum number of query embeddings to cache
        """
        super().__init__(
            model_name=model_name,
            k=k,
            alpha=alpha,
            verbose=verbose,
            device=device,
            local_files_only=local_files_only
        )

        # Add model_name attribute (not set by base class)
        self.model_name = model_name

        # Optimization: Query embedding LRU cache
        self.query_cache: OrderedDict[str, np.ndarray] = OrderedDict()
        self.query_cache_size = query_cache_size

        # Optimization: Pre-computed graph adjacency (for faster neighbor queries)
        self.graph_adjacency: Optional[Dict] = None

        # Cache configuration (from optimized version)
        self.cache_dir = cache_dir or os.path.expanduser('~/projects/embedding_cache/maniscope')
        self.use_cache = use_cache

    def _compute_cache_key(self, docs: List[str]) -> str:
        """
        Compute hash-based cache key from documents and model name.

        Args:
            docs: List of document texts

        Returns:
            Cache key string combining model name and document hash
        """
        content = '|'.join(docs)
        doc_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
        # Replace slashes in model name for valid filenames
        safe_model_name = self.model_name.replace('/', '_')
        return f"{safe_model_name}_{doc_hash}"

    def _get_cache_path(self, cache_key: str) -> Path:
        """
        Get full path to cache file.

        Args:
            cache_key: Cache key from _compute_cache_key()

        Returns:
            Path object for cache file
        """
        cache_dir = Path(self.cache_dir)
        return cache_dir / f"{cache_key}.pkl"

    def _load_embeddings_from_cache(self, cache_key: str) -> Optional[np.ndarray]:
        """
        Load embeddings from disk cache if available using joblib for faster deserialization.

        Args:
            cache_key: Cache key from _compute_cache_key()

        Returns:
            Cached embeddings array or None if not found
        """
        if not self.use_cache:
            return None

        cache_path = self._get_cache_path(cache_key)
        if cache_path.exists():
            if self.verbose:
                print(f"Loading embeddings from cache: {cache_path}")
            try:
                return joblib.load(cache_path)
            except Exception as e:
                if self.verbose:
                    print(f"Failed to load cache: {e}")
                return None
        return None

    def _save_embeddings_to_cache(self, embeddings: np.ndarray, cache_key: str):
        """
        Save embeddings to disk cache using joblib for faster serialization.

        Args:
            embeddings: Document embeddings array
            cache_key: Cache key from _compute_cache_key()
        """
        if not self.use_cache:
            return

        cache_path = self._get_cache_path(cache_key)
        try:
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            if self.verbose:
                print(f"Saving embeddings to cache: {cache_path}")
            joblib.dump(embeddings, cache_path)
        except Exception as e:
            if self.verbose:
                print(f"Failed to save cache: {e}")

    def _get_cached_query_embedding(self, query: str) -> Optional[np.ndarray]:
        """
        Retrieve cached query embedding or None if not cached.

        Args:
            query: Query text

        Returns:
            Cached embedding or None
        """
        return self.query_cache.get(query)

    def _cache_query_embedding(self, query: str, embedding: np.ndarray):
        """
        Cache query embedding with LRU eviction.

        Args:
            query: Query text
            embedding: Query embedding array
        """
        if query in self.query_cache:
            # Move to end (most recent)
            self.query_cache.move_to_end(query)
        else:
            if len(self.query_cache) >= self.query_cache_size:
                # Evict least recently used
                self.query_cache.popitem(last=False)
            self.query_cache[query] = embedding

    def _encode_query(self, query: str) -> np.ndarray:
        """
        Encode query with caching optimization.

        Args:
            query: Query text

        Returns:
            Query embedding array
        """
        cached = self._get_cached_query_embedding(query)
        if cached is not None:
            if self.verbose:
                print(f"  Using cached query embedding for: {query[:30]}...")
            return cached

        if self.verbose:
            print(f"  Encoding query: {query[:30]}...")
        embedding = self.model.encode([query])
        self._cache_query_embedding(query, embedding)
        return embedding

    def fit(self, docs: List[str]):
        """
        Build the manifold graph from document corpus with optimizations.

        Embeddings are cached to disk for efficiency. Subsequent calls with
        the same documents will load from cache instead of re-encoding.

        Args:
            docs: List of document texts

        Returns:
            self for method chaining
        """
        self.documents = docs

        # Try to load embeddings from cache
        cache_key = self._compute_cache_key(docs)
        cached_embeddings = self._load_embeddings_from_cache(cache_key)

        if cached_embeddings is not None:
            if self.verbose:
                print(f"Using cached embeddings for {len(docs)} documents")
            self.embeddings = cached_embeddings
        else:
            if self.verbose:
                print(f"Encoding {len(docs)} documents...")
            self.embeddings = self.model.encode(docs, show_progress_bar=self.verbose)
            self._save_embeddings_to_cache(self.embeddings, cache_key)

        if self.verbose:
            print(f"Building k-NN manifold graph (k={self.k})...")

        # Build k-NN manifold graph using cosine metric
        nbrs = NearestNeighbors(
            n_neighbors=min(self.k + 1, len(docs)),
            metric='cosine'
        ).fit(self.embeddings)
        distances, indices = nbrs.kneighbors(self.embeddings)

        self.G = nx.Graph()
        for i in range(len(docs)):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    self.G.add_edge(i, neighbor_idx, weight=dist)

        # Optimization: Pre-compute adjacency list for faster neighbor queries
        self.graph_adjacency = {
            node: [(neighbor, self.G[node][neighbor]['weight'])
                   for neighbor in self.G.neighbors(node)]
            for node in self.G.nodes()
        }

        if self.verbose:
            print(f"Graph built: {self.G.number_of_nodes()} nodes, "
                  f"{self.G.number_of_edges()} edges")
            print(f"Connected components: "
                  f"{nx.number_connected_components(self.G)}")

    def search_baseline(
        self,
        query: str,
        top_n: int = 5
    ) -> List[Tuple[str, float, int]]:
        """
        Baseline search using pure cosine similarity (telescope only).

        Args:
            query: Query text
            top_n: Number of results to return

        Returns:
            List of (document, score, doc_index) tuples sorted by score
        """
        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        q_emb = self._encode_query(query)
        sims = cosine_similarity(q_emb, self.embeddings)[0]
        idx = np.argsort(sims)[::-1][:top_n]
        return [(self.documents[i], float(sims[i]), int(i)) for i in idx]

    def search(
        self,
        query: str,
        top_n: int = 5,
        coarse_multiplier: int = 3
    ) -> List[Tuple[str, float, int]]:
        """
        Optimized Maniscope search: Telescope + Microscope (geodesic reranking).

        Incorporates batch shortest path computation, query caching, and early termination.

        Args:
            query: Query text
            top_n: Number of final results to return
            coarse_multiplier: Retrieve top_n * coarse_multiplier for reranking

        Returns:
            List of (document, score, doc_index) tuples sorted by hybrid score
        """
        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        # Phase 1: TELESCOPE - Broad retrieval with cosine similarity
        q_emb = self._encode_query(query)
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        # Retrieve broader set for reranking
        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        if self.verbose:
            print(f"\nTelescope: Retrieved top {coarse_size} candidates")

        # Phase 2: MICROSCOPE - Geodesic reranking on manifold
        # Select anchor node from coarse set (closest to query)
        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        if self.verbose:
            print(f"Microscope: Anchor node = {anchor_node} "
                  f"(cosine sim = {cosine_sims[anchor_node]:.3f})")

        # Optimization: Batch compute all geodesic distances from anchor
        try:
            all_geo_dists = nx.single_source_dijkstra_path_length(
                self.G, source=anchor_node, weight='weight'
            )
        except nx.NetworkXError:
            all_geo_dists = {}

        # Optimization: Vectorized computation with heap for early termination
        heap = []
        for i in coarse_idx:
            cosine_score = cosine_sims[i]
            geo_dist = all_geo_dists.get(i, None)
            if geo_dist is not None:
                geo_score = 1.0 / (1.0 + geo_dist)
            else:
                geo_score = 0.0  # Disconnected

            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            # Use heap to maintain top_n (min-heap, so use negative score)
            heappush(heap, (-final_score, i))
            if len(heap) > top_n:
                heappop(heap)

        # Extract results from heap
        results = []
        while heap:
            neg_score, i = heappop(heap)
            score = -neg_score
            results.append((self.documents[i], score, i))

        # Heap gives ascending order, reverse to descending
        results.reverse()

        return results

    def search_detailed(
        self,
        query: str,
        top_n: int = 5,
        coarse_multiplier: int = 3
    ) -> List[Dict]:
        """
        Optimized Maniscope search with detailed scoring information for analysis.

        Uses batch geodesic computation for efficiency.

        Args:
            query: Query text
            top_n: Number of results to return
            coarse_multiplier: Retrieve top_n * coarse_multiplier for reranking

        Returns:
            List of dicts with document, scores, and metadata:
            - doc_id: Document index
            - document: Document text
            - final_score: Hybrid score (alpha * cosine + (1-alpha) * geodesic)
            - cosine_score: Cosine similarity score
            - geo_score: Geodesic similarity score
            - geo_distance: Geodesic distance (None if disconnected)
            - connected: Whether node is connected to anchor in graph
        """
        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        q_emb = self._encode_query(query)
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        # Optimization: Batch compute all geodesic distances from anchor
        try:
            all_geo_dists = nx.single_source_dijkstra_path_length(
                self.G, source=anchor_node, weight='weight'
            )
        except nx.NetworkXError:
            all_geo_dists = {}

        results = []
        for i in coarse_idx:
            cosine_score = float(cosine_sims[i])
            geo_dist = all_geo_dists.get(i)
            if geo_dist is not None:
                geo_score = float(1.0 / (1.0 + geo_dist))
                connected = True
            else:
                geo_score = 0.0
                connected = False

            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            results.append({
                'doc_id': int(i),
                'document': self.documents[i],
                'final_score': float(final_score),
                'cosine_score': cosine_score,
                'geo_score': geo_score,
                'geo_distance': geo_dist,
                'connected': connected
            })

        results.sort(key=lambda x: x['final_score'], reverse=True)
        return results[:top_n]

    # Compatibility aliases for wrapper functions
    def search_maniscope(self, query: str, top_n: int = 5,
                         coarse_multiplier: int = 3) -> List[Tuple[str, float, int]]:
        """Alias for search() to maintain compatibility with wrapper functions."""
        return self.search(query, top_n, coarse_multiplier)

    def search_maniscope_detailed(self, query: str, top_n: int = 5,
                                  coarse_multiplier: int = 3) -> List[Dict]:
        """Alias for search_detailed() to maintain compatibility with wrapper functions."""
        return self.search_detailed(query, top_n, coarse_multiplier)


# =============================================================================
# Optimization v1: GPU + Graph Caching (3x faster)
# =============================================================================

class ManiscopeEngine_v1(ManiscopeEngine):
    """
    v1 Optimization: GPU Auto-detection + Graph Caching

    Target: 3× speedup (115ms → 40ms)

    Key improvements over v0 baseline:
    - GPU auto-detection for embeddings (3.5× faster encoding)
    - Graph caching between queries (eliminates 42ms rebuild)

    Performance: 40-45ms avg on subsequent queries
    Accuracy: MRR=1.0000 (same as v0, no regression)

    Note: Use with run_maniscope_reranker_v1() which implements graph caching
    """

    def __init__(self, model_name: str = 'all-MiniLM-L6-v2', k: int = 5,
                 alpha: float = 0.5, verbose: bool = False, device: Optional[str] = None,
                 local_files_only: bool = True):
        """
        Initialize v1 ManiscopeEngine with GPU auto-detection.

        Args:
            model_name: Sentence transformer model name
            k: Number of nearest neighbors for manifold graph
            alpha: Hybrid scoring weight (0=pure geodesic, 1=pure cosine)
            verbose: Print debug information
            device: Device ('cpu', 'cuda', or None for auto-detect)
            local_files_only: Use only cached models
        """
        # v1 Optimization: GPU auto-detection
        import torch
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            if verbose:
                gpu_info = f" (GPU: {torch.cuda.get_device_name(0)})" if device == 'cuda' else ""
                print(f"[v1] Auto-detected device: {device}{gpu_info}")

        # Initialize base class with detected device
        super().__init__(
            model_name=model_name,
            k=k,
            alpha=alpha,
            verbose=verbose,
            device=device,
            local_files_only=local_files_only
        )


# =============================================================================
# Optimization v2: Full Optimization (5x faster)
# =============================================================================

class ManiscopeEngine_v2(ManiscopeEngine):
    """
    v2 Optimization: Full Performance Optimization

    Target: 5× speedup (115ms → 20-25ms)

    Key improvements over v0 baseline:
    - All v1 optimizations (GPU + caching)
    - scipy.sparse for geodesic distances (4× faster than NetworkX)
    - FAISS for GPU-accelerated k-NN (4× faster than sklearn)
    - Vectorized hybrid scoring (4× faster than loops)

    Performance: 20-25ms avg
    Accuracy: MRR=1.0000 (same as v0/v1, no regression)
    """

    def __init__(self, model_name: str = 'all-MiniLM-L6-v2', k: int = 5,
                 alpha: float = 0.5, verbose: bool = False, device: Optional[str] = None,
                 local_files_only: bool = True, use_faiss: bool = True):
        """
        Initialize v2 ManiscopeEngine with full optimizations.

        Args:
            model_name: Sentence transformer model name
            k: Number of nearest neighbors for manifold graph
            alpha: Hybrid scoring weight (0=pure geodesic, 1=pure cosine)
            verbose: Print debug information
            device: Device ('cpu', 'cuda', or None for auto-detect)
            local_files_only: Use only cached models
            use_faiss: Use FAISS for k-NN (falls back to sklearn if unavailable)
        """
        # GPU auto-detection
        import torch
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            if verbose:
                gpu_info = f" (GPU: {torch.cuda.get_device_name(0)})" if device == 'cuda' else ""
                print(f"[v2] Auto-detected device: {device}{gpu_info}")

        super().__init__(
            model_name=model_name,
            k=k,
            alpha=alpha,
            verbose=verbose,
            device=device,
            local_files_only=local_files_only
        )

        self.use_faiss = use_faiss
        self.adj_matrix = None  # scipy sparse matrix for geodesic distances

        # Check FAISS availability
        if self.use_faiss:
            try:
                import faiss
                self.faiss_available = True
                if verbose:
                    print("[v2] FAISS available for GPU-accelerated k-NN")
            except ImportError:
                self.faiss_available = False
                if verbose:
                    print("[v2] FAISS not available, using sklearn fallback")
        else:
            self.faiss_available = False

    def fit(self, docs: List[str]):
        """
        Build manifold graph with v2 optimizations (scipy + FAISS).

        Args:
            docs: List of document texts
        """
        if self.verbose:
            print(f"[v2] Encoding {len(docs)} documents...")

        self.documents = docs
        self.embeddings = self.model.encode(docs, show_progress_bar=self.verbose)

        if self.verbose:
            print(f"[v2] Building k-NN manifold graph (k={self.k})...")

        # v2 Optimization: Use FAISS for k-NN if available
        if self.faiss_available and self.use_faiss:
            indices, distances = self._build_knn_faiss()
        else:
            indices, distances = self._build_knn_sklearn()

        # Build NetworkX graph (still needed for some operations)
        self.G = nx.Graph()
        for i in range(len(docs)):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    self.G.add_edge(i, neighbor_idx, weight=dist)

        # v2 Optimization: Build scipy sparse adjacency matrix for fast geodesic
        self._build_sparse_adjacency(indices, distances)

        if self.verbose:
            print(f"[v2] Graph built: {self.G.number_of_nodes()} nodes, "
                  f"{self.G.number_of_edges()} edges")
            print(f"[v2] Connected components: {nx.number_connected_components(self.G)}")

    def _build_knn_sklearn(self) -> Tuple[np.ndarray, np.ndarray]:
        """Build k-NN using sklearn (fallback)."""
        nbrs = NearestNeighbors(
            n_neighbors=min(self.k + 1, len(self.documents)),
            metric='cosine'
        ).fit(self.embeddings)
        distances, indices = nbrs.kneighbors(self.embeddings)
        return indices, distances

    def _build_knn_faiss(self) -> Tuple[np.ndarray, np.ndarray]:
        """Build k-NN using FAISS (v2 optimization)."""
        import faiss
        import torch

        # Normalize embeddings for cosine similarity (inner product on unit vectors)
        embeddings_norm = self.embeddings / np.linalg.norm(self.embeddings, axis=1, keepdims=True)
        embeddings_norm = embeddings_norm.astype('float32')

        d = embeddings_norm.shape[1]  # Embedding dimension
        k = min(self.k + 1, len(self.documents))

        # Build FAISS index (GPU if available)
        if torch.cuda.is_available() and hasattr(faiss, 'StandardGpuResources'):
            try:
                res = faiss.StandardGpuResources()
                index = faiss.GpuIndexFlatIP(res, d)
                if self.verbose:
                    print("[v2] Using GPU FAISS index")
            except Exception as e:
                if self.verbose:
                    print(f"[v2] GPU FAISS failed ({e}), using CPU")
                index = faiss.IndexFlatIP(d)
        else:
            index = faiss.IndexFlatIP(d)

        index.add(embeddings_norm)

        # Query k-NN (inner product = cosine for normalized vectors)
        # FAISS returns squared distances for IP, convert to cosine distance
        similarities, indices = index.search(embeddings_norm, k)
        distances = 1.0 - similarities  # Convert similarity to distance

        return indices, distances

    def _build_sparse_adjacency(self, indices: np.ndarray, distances: np.ndarray):
        """Build scipy sparse adjacency matrix for fast geodesic computation."""
        from scipy.sparse import csr_matrix

        n = len(self.documents)
        row, col, data = [], [], []

        for i in range(n):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    row.append(i)
                    col.append(neighbor_idx)
                    data.append(dist)

        self.adj_matrix = csr_matrix((data, (row, col)), shape=(n, n))

        if self.verbose:
            print(f"[v2] Built sparse adjacency matrix: {self.adj_matrix.shape}, "
                  f"{self.adj_matrix.nnz} non-zeros")

    def search_maniscope(self, query: str, top_n: int = 5,
                         coarse_multiplier: int = 3) -> List[Tuple[str, float, int]]:
        """
        v2 optimized Maniscope search with scipy geodesic computation.

        Args:
            query: Query text
            top_n: Number of results to return
            coarse_multiplier: Retrieve top_n * coarse_multiplier for reranking

        Returns:
            List of (document, score, doc_index) tuples
        """
        from scipy.sparse.csgraph import dijkstra

        # Phase 1: TELESCOPE
        q_emb = self.model.encode([query])
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        if self.verbose:
            print(f"\n[v2] Telescope: Retrieved top {coarse_size} candidates")

        # Phase 2: MICROSCOPE with v2 optimization (scipy dijkstra)
        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        if self.verbose:
            print(f"[v2] Microscope: Anchor node = {anchor_node} "
                  f"(cosine sim = {cosine_sims[anchor_node]:.3f})")

        # v2 Optimization: Vectorized geodesic computation with scipy
        geo_dists = dijkstra(
            self.adj_matrix,
            indices=[anchor_node],
            directed=False,
            return_predecessors=False
        )[0]

        # v2 Optimization: Vectorized hybrid scoring
        cosine_scores = cosine_sims[coarse_idx]
        geo_dists_batch = geo_dists[coarse_idx]

        # Handle disconnected nodes (infinite distance)
        geo_scores = np.where(
            np.isinf(geo_dists_batch),
            0.0,
            1.0 / (1.0 + geo_dists_batch)
        )

        # Hybrid scoring (vectorized)
        final_scores = self.alpha * cosine_scores + (1 - self.alpha) * geo_scores

        # Sort and return top_n
        sorted_idx = np.argsort(final_scores)[::-1][:top_n]
        results = [
            (self.documents[coarse_idx[i]], float(final_scores[i]), int(coarse_idx[i]))
            for i in sorted_idx
        ]

        return results

    def search_maniscope_detailed(self, query: str, top_n: int = 5,
                                  coarse_multiplier: int = 3) -> List[dict]:
        """
        v2 optimized detailed search with scoring breakdown.

        Args:
            query: Query text
            top_n: Number of results to return
            coarse_multiplier: Retrieve top_n * coarse_multiplier for reranking

        Returns:
            List of dicts with detailed scores
        """
        from scipy.sparse.csgraph import dijkstra

        q_emb = self.model.encode([query])
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        # Vectorized geodesic computation
        geo_dists = dijkstra(
            self.adj_matrix,
            indices=[anchor_node],
            directed=False,
            return_predecessors=False
        )[0]

        # Build results
        results = []
        for i in coarse_idx:
            cosine_score = float(cosine_sims[i])
            geo_dist = float(geo_dists[i]) if not np.isinf(geo_dists[i]) else None

            if geo_dist is not None:
                geo_score = float(1.0 / (1.0 + geo_dist))
                connected = True
            else:
                geo_score = 0.0
                connected = False

            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            results.append({
                'doc_id': int(i),
                'document': self.documents[i],
                'final_score': float(final_score),
                'cosine_score': cosine_score,
                'geo_score': geo_score,
                'geo_distance': geo_dist,
                'connected': connected
            })

        results.sort(key=lambda x: x['final_score'], reverse=True)
        return results[:top_n]


# =============================================================================
# Optimization v2o: Ultimate Optimization - Best of All Worlds (20-235x faster)
# =============================================================================

class ManiscopeEngine_v2o(ManiscopeEngine):
    """
    v2o Optimization: Ultimate Performance - Combines ALL optimizations

    Target: 20-235× speedup (115ms → 0.4-0.6ms with warm cache, 4-20ms cold)

    This is the RECOMMENDED version combining the best of v1, v2, and v3:

    From v1 (GPU):
    - GPU auto-detection for embeddings (3.5× faster encoding)
    - Smart device selection (CUDA if available, CPU fallback)

    From v2 (Algorithmic):
    - scipy.sparse for geodesic distances (4× faster than NetworkX)
    - FAISS for GPU-accelerated k-NN (4× faster than sklearn)
    - Vectorized hybrid scoring (4× faster than loops)
    - Batch geodesic computation

    From v3 (Caching):
    - Persistent disk cache for embeddings with joblib (survives restarts!)
    - Query embedding LRU cache (100 queries in memory)
    - Graph adjacency list pre-computation

    Performance Benchmarks (Real-World Results):
    - MS MARCO (200 queries): 132ms → 0.58ms (229× speedup)
    - TREC-COVID (50 queries): 85ms → 0.38ms (226× speedup)
    - SciFact: 92ms → 0.39ms (235× speedup)
    - Cold cache (first run): ~5-25× speedup (still faster than v0!)
    - Warm cache (cached embeddings): 20-235× speedup

    Accuracy: MRR=1.0000 (identical to v0, zero regression)

    Best for: Production deployments, repeated experiments, grid search
    Memory: Moderate (caches embeddings + query cache + sparse matrices)
    GPU: Optional but recommended (falls back to CPU gracefully)
    """

    def __init__(
        self,
        model_name: str = 'all-MiniLM-L6-v2',
        k: int = 5,
        alpha: float = 0.5,
        verbose: bool = False,
        device: Optional[str] = None,  # Auto-detect GPU
        local_files_only: bool = True,
        cache_dir: Optional[str] = None,
        use_cache: bool = True,
        query_cache_size: int = 100,
        use_faiss: bool = True
    ):
        """
        Initialize the ultimate optimized ManiscopeEngine.

        Args:
            model_name: Sentence transformer model name
            k: Number of nearest neighbors for manifold graph construction
            alpha: Weight for hybrid scoring (0=pure geodesic, 1=pure cosine)
            verbose: Print debug information
            device: Device ('cpu', 'cuda', or None for auto-detect)
            local_files_only: Use only cached models without checking HuggingFace
            cache_dir: Directory for embedding cache
            use_cache: Enable disk-based caching of embeddings
            query_cache_size: Maximum number of query embeddings to cache
            use_faiss: Use FAISS for k-NN (falls back to sklearn if unavailable)
        """
        # v1 Optimization: GPU auto-detection
        import torch
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            if verbose:
                gpu_info = f" (GPU: {torch.cuda.get_device_name(0)})" if device == 'cuda' else ""
                print(f"[v2o] Auto-detected device: {device}{gpu_info}")

        # Initialize base class
        super().__init__(
            model_name=model_name,
            k=k,
            alpha=alpha,
            verbose=verbose,
            device=device,
            local_files_only=local_files_only
        )

        # Add model_name attribute
        self.model_name = model_name

        # v3 Optimization: Query embedding LRU cache
        self.query_cache: OrderedDict[str, np.ndarray] = OrderedDict()
        self.query_cache_size = query_cache_size

        # v3 Optimization: Cache configuration
        self.cache_dir = cache_dir or os.path.expanduser('~/projects/embedding_cache/maniscope')
        self.use_cache = use_cache

        # v2 Optimization: FAISS configuration
        self.use_faiss = use_faiss
        self.adj_matrix = None  # scipy sparse matrix for geodesic distances

        # Check FAISS availability
        if self.use_faiss:
            try:
                import faiss
                self.faiss_available = True
                if verbose:
                    print("[v2o] FAISS available for GPU-accelerated k-NN")
            except ImportError:
                self.faiss_available = False
                if verbose:
                    print("[v2o] FAISS not available, using sklearn fallback")
        else:
            self.faiss_available = False

    # =============================================================================
    # v3 Cache Management Methods
    # =============================================================================

    def _compute_cache_key(self, docs: List[str]) -> str:
        """Compute hash-based cache key from documents and model name."""
        content = '|'.join(docs)
        doc_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
        safe_model_name = self.model_name.replace('/', '_')
        return f"{safe_model_name}_{doc_hash}"

    def _get_cache_path(self, cache_key: str) -> Path:
        """Get full path to cache file."""
        cache_dir = Path(self.cache_dir)
        return cache_dir / f"{cache_key}.pkl"

    def _load_embeddings_from_cache(self, cache_key: str) -> Optional[np.ndarray]:
        """Load embeddings from disk cache using joblib."""
        if not self.use_cache:
            return None

        cache_path = self._get_cache_path(cache_key)
        if cache_path.exists():
            if self.verbose:
                print(f"[v2o] Loading embeddings from cache: {cache_path.name}")
            try:
                return joblib.load(cache_path)
            except Exception as e:
                if self.verbose:
                    print(f"[v2o] Failed to load cache: {e}")
                return None
        return None

    def _save_embeddings_to_cache(self, embeddings: np.ndarray, cache_key: str):
        """Save embeddings to disk cache using joblib."""
        if not self.use_cache:
            return

        cache_path = self._get_cache_path(cache_key)
        try:
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            if self.verbose:
                print(f"[v2o] Saving embeddings to cache: {cache_path.name}")
            joblib.dump(embeddings, cache_path)
        except Exception as e:
            if self.verbose:
                print(f"[v2o] Failed to save cache: {e}")

    def _get_cached_query_embedding(self, query: str) -> Optional[np.ndarray]:
        """Retrieve cached query embedding or None if not cached."""
        return self.query_cache.get(query)

    def _cache_query_embedding(self, query: str, embedding: np.ndarray):
        """Cache query embedding with LRU eviction."""
        if query in self.query_cache:
            self.query_cache.move_to_end(query)
        else:
            if len(self.query_cache) >= self.query_cache_size:
                self.query_cache.popitem(last=False)
            self.query_cache[query] = embedding

    def _encode_query(self, query: str) -> np.ndarray:
        """Encode query with caching optimization."""
        cached = self._get_cached_query_embedding(query)
        if cached is not None:
            if self.verbose:
                print(f"[v2o] Using cached query embedding")
            return cached

        embedding = self.model.encode([query])
        self._cache_query_embedding(query, embedding)
        return embedding

    # =============================================================================
    # v2 FAISS k-NN Methods
    # =============================================================================

    def _build_knn_sklearn(self) -> Tuple[np.ndarray, np.ndarray]:
        """Build k-NN using sklearn (fallback)."""
        nbrs = NearestNeighbors(
            n_neighbors=min(self.k + 1, len(self.documents)),
            metric='cosine'
        ).fit(self.embeddings)
        distances, indices = nbrs.kneighbors(self.embeddings)
        return indices, distances

    def _build_knn_faiss(self) -> Tuple[np.ndarray, np.ndarray]:
        """Build k-NN using FAISS (v2 optimization)."""
        import faiss
        import torch

        # Normalize embeddings for cosine similarity
        embeddings_norm = self.embeddings / np.linalg.norm(self.embeddings, axis=1, keepdims=True)
        embeddings_norm = embeddings_norm.astype('float32')

        d = embeddings_norm.shape[1]
        k = min(self.k + 1, len(self.documents))

        # Build FAISS index (GPU if available)
        if torch.cuda.is_available() and hasattr(faiss, 'StandardGpuResources'):
            try:
                res = faiss.StandardGpuResources()
                index = faiss.GpuIndexFlatIP(res, d)
                if self.verbose:
                    print("[v2o] Using GPU FAISS index")
            except Exception as e:
                if self.verbose:
                    print(f"[v2o] GPU FAISS failed ({e}), using CPU")
                index = faiss.IndexFlatIP(d)
        else:
            index = faiss.IndexFlatIP(d)

        index.add(embeddings_norm)
        similarities, indices = index.search(embeddings_norm, k)
        distances = 1.0 - similarities

        return indices, distances

    def _build_sparse_adjacency(self, indices: np.ndarray, distances: np.ndarray):
        """Build scipy sparse adjacency matrix for fast geodesic computation."""
        from scipy.sparse import csr_matrix

        n = len(self.documents)
        row, col, data = [], [], []

        for i in range(n):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    row.append(i)
                    col.append(neighbor_idx)
                    data.append(dist)

        self.adj_matrix = csr_matrix((data, (row, col)), shape=(n, n))

        if self.verbose:
            print(f"[v2o] Built sparse adjacency matrix: {self.adj_matrix.shape}, "
                  f"{self.adj_matrix.nnz} non-zeros")

    # =============================================================================
    # Core Methods
    # =============================================================================

    def fit(self, docs: List[str]):
        """
        Build the manifold graph with ALL optimizations.

        Combines v2 FAISS + v3 caching for maximum performance.
        """
        self.documents = docs

        # v3 Optimization: Try to load embeddings from cache
        cache_key = self._compute_cache_key(docs)
        cached_embeddings = self._load_embeddings_from_cache(cache_key)

        if cached_embeddings is not None:
            if self.verbose:
                print(f"[v2o] Using cached embeddings for {len(docs)} documents")
            self.embeddings = cached_embeddings
        else:
            if self.verbose:
                print(f"[v2o] Encoding {len(docs)} documents...")
            self.embeddings = self.model.encode(docs, show_progress_bar=self.verbose)
            self._save_embeddings_to_cache(self.embeddings, cache_key)

        if self.verbose:
            print(f"[v2o] Building k-NN manifold graph (k={self.k})...")

        # v2 Optimization: Use FAISS for k-NN if available
        if self.faiss_available and self.use_faiss:
            indices, distances = self._build_knn_faiss()
        else:
            indices, distances = self._build_knn_sklearn()

        # Build NetworkX graph (needed for some operations)
        self.G = nx.Graph()
        for i in range(len(docs)):
            for neighbor_idx, dist in zip(indices[i], distances[i]):
                if i != neighbor_idx:
                    self.G.add_edge(i, neighbor_idx, weight=dist)

        # v2 Optimization: Build scipy sparse adjacency matrix
        self._build_sparse_adjacency(indices, distances)

        if self.verbose:
            print(f"[v2o] Graph built: {self.G.number_of_nodes()} nodes, "
                  f"{self.G.number_of_edges()} edges")
            print(f"[v2o] Connected components: {nx.number_connected_components(self.G)}")

        return self

    def search_baseline(
        self,
        query: str,
        top_n: int = 5
    ) -> List[Tuple[str, float, int]]:
        """Baseline search using pure cosine similarity."""
        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        q_emb = self._encode_query(query)
        sims = cosine_similarity(q_emb, self.embeddings)[0]
        idx = np.argsort(sims)[::-1][:top_n]
        return [(self.documents[i], float(sims[i]), int(i)) for i in idx]

    def search(
        self,
        query: str,
        top_n: int = 5,
        coarse_multiplier: int = 3
    ) -> List[Tuple[str, float, int]]:
        """
        Ultimate optimized Maniscope search.

        Combines v2 scipy dijkstra + v3 query caching for maximum speed.
        """
        from scipy.sparse.csgraph import dijkstra

        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        # Phase 1: TELESCOPE - Broad retrieval with cosine similarity
        q_emb = self._encode_query(query)
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        if self.verbose:
            print(f"[v2o] Telescope: Retrieved top {coarse_size} candidates")

        # Phase 2: MICROSCOPE - Geodesic reranking
        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        if self.verbose:
            print(f"[v2o] Microscope: Anchor node = {anchor_node} "
                  f"(cosine sim = {cosine_sims[anchor_node]:.3f})")

        # v2 Optimization: Vectorized geodesic computation with scipy
        geo_dists = dijkstra(
            self.adj_matrix,
            indices=[anchor_node],
            directed=False,
            return_predecessors=False
        )[0]

        # v2 Optimization: Vectorized hybrid scoring
        cosine_scores = cosine_sims[coarse_idx]
        geo_dists_batch = geo_dists[coarse_idx]

        # Handle disconnected nodes
        geo_scores = np.where(
            np.isinf(geo_dists_batch),
            0.0,
            1.0 / (1.0 + geo_dists_batch)
        )

        # Hybrid scoring (vectorized)
        final_scores = self.alpha * cosine_scores + (1 - self.alpha) * geo_scores

        # Sort and return top_n
        sorted_idx = np.argsort(final_scores)[::-1][:top_n]
        results = [
            (self.documents[coarse_idx[i]], float(final_scores[i]), int(coarse_idx[i]))
            for i in sorted_idx
        ]

        return results

    def search_detailed(
        self,
        query: str,
        top_n: int = 5,
        coarse_multiplier: int = 3
    ) -> List[Dict]:
        """
        Ultimate optimized detailed search with scoring breakdown.
        """
        from scipy.sparse.csgraph import dijkstra

        if self.embeddings is None:
            raise ValueError("Must call fit() before search")

        q_emb = self._encode_query(query)
        cosine_sims = cosine_similarity(q_emb, self.embeddings)[0]

        coarse_size = min(top_n * coarse_multiplier, len(self.documents))
        coarse_idx = np.argsort(cosine_sims)[::-1][:coarse_size]

        anchor_candidates = coarse_idx[:min(5, len(coarse_idx))]
        anchor_node = anchor_candidates[0]

        # Vectorized geodesic computation
        geo_dists = dijkstra(
            self.adj_matrix,
            indices=[anchor_node],
            directed=False,
            return_predecessors=False
        )[0]

        # Build results
        results = []
        for i in coarse_idx:
            cosine_score = float(cosine_sims[i])
            geo_dist = float(geo_dists[i]) if not np.isinf(geo_dists[i]) else None

            if geo_dist is not None:
                geo_score = float(1.0 / (1.0 + geo_dist))
                connected = True
            else:
                geo_score = 0.0
                connected = False

            final_score = self.alpha * cosine_score + (1 - self.alpha) * geo_score

            results.append({
                'doc_id': int(i),
                'document': self.documents[i],
                'final_score': float(final_score),
                'cosine_score': cosine_score,
                'geo_score': geo_score,
                'geo_distance': geo_dist,
                'connected': connected
            })

        results.sort(key=lambda x: x['final_score'], reverse=True)
        return results[:top_n]

    # Compatibility aliases
    def search_maniscope(self, query: str, top_n: int = 5,
                         coarse_multiplier: int = 3) -> List[Tuple[str, float, int]]:
        """Alias for search() to maintain compatibility."""
        return self.search(query, top_n, coarse_multiplier)

    def search_maniscope_detailed(self, query: str, top_n: int = 5,
                                  coarse_multiplier: int = 3) -> List[Dict]:
        """Alias for search_detailed() to maintain compatibility."""
        return self.search_detailed(query, top_n, coarse_multiplier)


# =============================================================================
# Utility Functions
# =============================================================================

def compare_maniscope_performance(
    engine1: ManiscopeEngine,
    engine2: ManiscopeEngine,
    query: str,
    top_n: int = 5,
    num_runs: int = 10,
    coarse_multiplier: int = 3
) -> Dict[str, Any]:
    """
    Compare performance of two ManiscopeEngine instances (any version: v0, v1, v2, v3).

    This function benchmarks two engines by running the same query multiple times
    and comparing their execution times and result consistency.

    Args:
        engine1: First engine instance (any version)
        engine2: Second engine instance (any version)
        query: Query text to benchmark
        top_n: Number of results to return
        num_runs: Number of benchmark runs for averaging
        coarse_multiplier: Coarse retrieval multiplier (passed to search)

    Returns:
        Dict with performance comparison:
        - engine1_name: Name of first engine (e.g., "ManiscopeEngine_v1")
        - engine2_name: Name of second engine (e.g., "ManiscopeEngine_v2")
        - engine1_time: Total time for engine1 (seconds)
        - engine2_time: Total time for engine2 (seconds)
        - engine1_avg: Average time per query for engine1 (ms)
        - engine2_avg: Average time per query for engine2 (ms)
        - speedup: engine1_time / engine2_time
        - results_consistent: Whether both engines returned same results
        - query: The query that was benchmarked
        - top_n: Number of results requested
        - num_runs: Number of runs performed

    Example:
        >>> from maniscope_engine import ManiscopeEngine, ManiscopeEngine_v1
        >>> engine_v0 = ManiscopeEngine(k=5, alpha=0.3)
        >>> engine_v1 = ManiscopeEngine_v1(k=5, alpha=0.3)
        >>> docs = ["document 1", "document 2", ...]
        >>> engine_v0.fit(docs)
        >>> engine_v1.fit(docs)
        >>> results = compare_maniscope_performance(engine_v0, engine_v1, "query", num_runs=10)
        >>> print(f"Speedup: {results['speedup']:.2f}x")
    """
    import time

    # Benchmark engine1
    start_time = time.time()
    results1 = []
    for _ in range(num_runs):
        result = engine1.search_maniscope(query, top_n=top_n, coarse_multiplier=coarse_multiplier)
        results1.append(result)
    time1 = time.time() - start_time

    # Benchmark engine2
    start_time = time.time()
    results2 = []
    for _ in range(num_runs):
        result = engine2.search_maniscope(query, top_n=top_n, coarse_multiplier=coarse_multiplier)
        results2.append(result)
    time2 = time.time() - start_time

    # Check consistency (compare document IDs and scores)
    consistent = True
    if len(results1) > 0 and len(results2) > 0:
        # Compare first run results
        r1 = results1[0]
        r2 = results2[0]
        if len(r1) == len(r2):
            for (doc1, score1, idx1), (doc2, score2, idx2) in zip(r1, r2):
                # Check if document indices match and scores are close
                if idx1 != idx2 or abs(score1 - score2) > 1e-6:
                    consistent = False
                    break
        else:
            consistent = False

    return {
        'engine1_name': engine1.__class__.__name__,
        'engine2_name': engine2.__class__.__name__,
        'engine1_time': time1,
        'engine2_time': time2,
        'engine1_avg': (time1 / num_runs) * 1000,  # Convert to ms
        'engine2_avg': (time2 / num_runs) * 1000,  # Convert to ms
        'speedup': time1 / time2 if time2 > 0 else float('inf'),
        'results_consistent': consistent,
        'query': query,
        'top_n': top_n,
        'num_runs': num_runs
    }
