src/data/loader.py

"""
WikiData singleton for accessing Wikipedia embeddings, link graph, and titles.

Usage:
    from src.data.loader import wiki_data

    # First access triggers lazy loading (~30-60 seconds)
    wiki_data.get_title(1234)
    wiki_data.get_embedding("Albert Einstein")
    wiki_data.get_links("Python (programming language)")
    wiki_data.similarity("Physics", "Mathematics")
"""

from __future__ import annotations

import json
import logging
from functools import lru_cache

import faiss
import msgpack
import numpy as np

from src.config import (
    DATA_DIR,
    EMBEDDING_DIM,
    EMBEDDINGS_PATH,
    LINK_GRAPH_PATH,
    TITLE_TO_IDX_PATH,
    TITLES_PATH,
)

# Path for persisted FAISS index (avoids rebuilding on each run)
FAISS_INDEX_PATH = DATA_DIR / "faiss_index.bin"

logger = logging.getLogger(__name__)


class WikiData:
    """
    Lazy-loading singleton for Wikipedia data.

    Loads data on first access to any method. All data is shared across
    the application via the module-level `wiki_data` instance.

    Attributes:
        titles: List of all Wikipedia article titles (indexed)
        title_to_idx: Dict mapping title string to index
        link_graph: Dict mapping source index to list of target indices
        embeddings: Memory-mapped numpy array of embeddings
        faiss_index: FAISS index for fast similarity search
    """

    _instance: WikiData | None = None

    def __new__(cls) -> WikiData:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def _ensure_loaded(self) -> None:
        """Load all data on first access."""
        if self._initialized:
            return

        logger.info("Loading Wikipedia data (this may take a minute)...")

        self._load_titles()
        self._load_title_to_idx()
        self._load_link_graph()
        self._load_embeddings()
        self._build_faiss_index()

        self._initialized = True
        logger.info("Wikipedia data loaded successfully")

    def _load_titles(self) -> None:
        """Load titles.json into list."""
        logger.info(f"Loading titles from {TITLES_PATH}...")
        with open(TITLES_PATH, encoding="utf-8") as f:
            self._titles: list[str] = json.load(f)
        logger.info(f"Loaded {len(self._titles):,} titles")

    def _load_title_to_idx(self) -> None:
        """Load title_to_idx.json into dict."""
        logger.info(f"Loading title->index mapping from {TITLE_TO_IDX_PATH}...")
        with open(TITLE_TO_IDX_PATH, encoding="utf-8") as f:
            self._title_to_idx: dict[str, int] = json.load(f)
        logger.info(f"Loaded {len(self._title_to_idx):,} mappings")

    def _load_link_graph(self) -> None:
        """Load link_graph.msgpack into dict."""
        logger.info(f"Loading link graph from {LINK_GRAPH_PATH}...")
        with open(LINK_GRAPH_PATH, "rb") as f:
            self._link_graph: dict[int, list[int]] = msgpack.load(f)
        logger.info(f"Loaded {len(self._link_graph):,} articles with outgoing links")

    def _load_embeddings(self) -> None:
        """Memory-map embeddings.npy for efficient access."""
        logger.info(f"Memory-mapping embeddings from {EMBEDDINGS_PATH}...")
        self._embeddings: np.ndarray = np.load(EMBEDDINGS_PATH, mmap_mode="r")
        logger.info(f"Embeddings shape: {self._embeddings.shape}, dtype: {self._embeddings.dtype}")

    def _build_faiss_index(self) -> None:
        """Load or build FAISS index (memory-efficient for 32GB RAM systems)."""
        # Try to load pre-built index from disk
        if FAISS_INDEX_PATH.exists():
            logger.info(f"Loading pre-built FAISS index from {FAISS_INDEX_PATH}...")
            self._faiss_index = faiss.read_index(str(FAISS_INDEX_PATH))
            logger.info(f"Loaded FAISS index with {self._faiss_index.ntotal:,} vectors")
            return

        # Build IVF index (approximate but memory-efficient)
        # IndexIVFFlat uses centroids + inverted lists, much lower memory than IndexFlatIP
        logger.info("Building FAISS IVF index (this may take a while)...")

        total = len(self._embeddings)
        nlist = 4096  # Number of clusters (sqrt(n) is typical)

        # Create IVF index with inner product (for cosine similarity on normalized vectors)
        quantizer = faiss.IndexFlatIP(EMBEDDING_DIM)
        self._faiss_index = faiss.IndexIVFFlat(quantizer, EMBEDDING_DIM, nlist, faiss.METRIC_INNER_PRODUCT)

        # Train on a sample (IVF requires training)
        logger.info("  Training index on sample...")
        sample_size = min(500_000, total)
        sample_indices = np.random.choice(total, sample_size, replace=False)
        sample = self._embeddings[sample_indices].astype(np.float32)
        faiss.normalize_L2(sample)
        self._faiss_index.train(sample)
        logger.info("  Training complete")

        # Add vectors in batches
        batch_size = 100_000
        for start_idx in range(0, total, batch_size):
            end_idx = min(start_idx + batch_size, total)

            batch = self._embeddings[start_idx:end_idx].astype(np.float32)
            faiss.normalize_L2(batch)
            self._faiss_index.add(batch)

            if (start_idx // batch_size) % 10 == 0:
                logger.info(f"  Added {end_idx:,}/{total:,} vectors ({100*end_idx/total:.1f}%)")

        # Set search parameters (higher nprobe = more accurate but slower)
        self._faiss_index.nprobe = 64

        # Save to disk for future runs
        logger.info(f"Saving FAISS index to {FAISS_INDEX_PATH}...")
        faiss.write_index(self._faiss_index, str(FAISS_INDEX_PATH))

        logger.info(f"FAISS index built with {self._faiss_index.ntotal:,} vectors")

    # =========================================================================
    # Core Accessors
    # =========================================================================

    def get_title(self, idx: int) -> str:
        """Get article title by index."""
        self._ensure_loaded()
        if 0 <= idx < len(self._titles):
            return self._titles[idx]
        raise IndexError(f"Index {idx} out of range [0, {len(self._titles)})")

    def get_index(self, title: str) -> int | None:
        """Get index for article title, or None if not found."""
        self._ensure_loaded()
        return self._title_to_idx.get(title)

    def has_article(self, title: str) -> bool:
        """Check if article exists in the dataset."""
        self._ensure_loaded()
        return title in self._title_to_idx

    def is_traversable(self, title: str) -> bool:
        """Check if article has outgoing links (can be navigated from)."""
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return False
        return idx in self._link_graph

    def article_count(self) -> int:
        """Total number of articles with embeddings."""
        self._ensure_loaded()
        return len(self._titles)

    def traversable_count(self) -> int:
        """Number of articles with outgoing links."""
        self._ensure_loaded()
        return len(self._link_graph)

    # =========================================================================
    # Embedding Accessors
    # =========================================================================

    def get_embedding(self, title: str) -> np.ndarray | None:
        """Get embedding for article by title, or None if not found."""
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return None
        return self._embeddings[idx].astype(np.float32)

    def get_embedding_by_idx(self, idx: int) -> np.ndarray:
        """Get embedding by index (faster, no lookup)."""
        self._ensure_loaded()
        return self._embeddings[idx].astype(np.float32)

    def get_normalized_embedding(self, title: str) -> np.ndarray | None:
        """Get L2-normalized embedding for article (for cosine similarity)."""
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return None
        return self._normalize_embedding(self._embeddings[idx])

    def _normalize_embedding(self, emb: np.ndarray) -> np.ndarray:
        """Normalize a single embedding to unit length (for cosine similarity)."""
        emb_f32 = emb.astype(np.float32)
        norm = np.linalg.norm(emb_f32)
        if norm > 0:
            emb_f32 /= norm
        return emb_f32

    def get_embeddings_batch(self, titles: list[str]) -> np.ndarray:
        """Get embeddings for multiple titles as a batch."""
        self._ensure_loaded()
        indices = [self._title_to_idx.get(t) for t in titles]
        valid_indices = [i for i in indices if i is not None]
        if not valid_indices:
            return np.array([])
        return self._embeddings[valid_indices].astype(np.float32)

    # =========================================================================
    # Similarity Functions (FAISS-accelerated)
    # =========================================================================

    def similarity(self, title_a: str, title_b: str) -> float | None:
        """
        Compute cosine similarity between two articles.

        Returns None if either article is not found.
        """
        self._ensure_loaded()
        idx_a = self._title_to_idx.get(title_a)
        idx_b = self._title_to_idx.get(title_b)
        if idx_a is None or idx_b is None:
            return None

        # Normalize embeddings on-demand for cosine similarity
        emb_a = self._normalize_embedding(self._embeddings[idx_a])
        emb_b = self._normalize_embedding(self._embeddings[idx_b])
        return float(np.dot(emb_a, emb_b))

    def rank_by_similarity(
        self, candidates: list[str], target: str
    ) -> list[tuple[str, float]]:
        """
        Rank candidate articles by similarity to target.

        Returns list of (title, similarity) tuples, sorted by similarity descending.
        """
        self._ensure_loaded()
        target_idx = self._title_to_idx.get(target)
        if target_idx is None:
            return []

        target_emb = self._normalize_embedding(self._embeddings[target_idx]).reshape(1, -1)

        # Get candidate indices
        candidate_indices = []
        valid_candidates = []
        for c in candidates:
            idx = self._title_to_idx.get(c)
            if idx is not None:
                candidate_indices.append(idx)
                valid_candidates.append(c)

        if not candidate_indices:
            return []

        # Compute similarities - normalize candidates on-demand
        candidate_embs = np.array([
            self._normalize_embedding(self._embeddings[i])
            for i in candidate_indices
        ])
        similarities = np.dot(candidate_embs, target_emb.T).flatten()

        # Sort by similarity descending
        ranked = sorted(
            zip(valid_candidates, similarities, strict=True),
            key=lambda x: x[1],
            reverse=True,
        )
        return [(title, float(sim)) for title, sim in ranked]

    def nearest_neighbors(self, title: str, k: int = 10) -> list[tuple[str, float]]:
        """
        Find k nearest neighbors to an article using FAISS.

        Returns list of (title, similarity) tuples.
        """
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return []

        query = self._normalize_embedding(self._embeddings[idx]).reshape(1, -1)
        similarities, indices = self._faiss_index.search(query, k + 1)

        # Skip the first result (the article itself)
        results = []
        for sim, neighbor_idx in zip(similarities[0], indices[0], strict=True):
            if neighbor_idx != idx and neighbor_idx >= 0:
                results.append((self._titles[neighbor_idx], float(sim)))
            if len(results) >= k:
                break

        return results

    # =========================================================================
    # Graph Accessors
    # =========================================================================

    def get_links(self, title: str) -> list[str]:
        """Get outgoing links from an article as titles."""
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None or idx not in self._link_graph:
            return []

        link_indices = self._link_graph[idx]
        return [self._titles[i] for i in link_indices if 0 <= i < len(self._titles)]

    def get_links_by_idx(self, idx: int) -> list[int]:
        """Get outgoing link indices for an article."""
        self._ensure_loaded()
        return self._link_graph.get(idx, [])

    @lru_cache(maxsize=10000)  # noqa: B019 - singleton pattern mitigates leak
    def get_popularity(self, title: str) -> int:
        """
        Get popularity score (number of inbound links) for an article.

        Computed on-demand and cached. First call for a new title requires
        iterating through the link graph, which may be slow.
        """
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return 0

        # Count how many articles link to this one
        count = 0
        for links in self._link_graph.values():
            if idx in links:
                count += 1

        return count

    def get_inbound_links(self, title: str) -> list[str]:
        """
        Get all articles that link to this article.

        Warning: This is slow (O(n) where n = number of articles).
        """
        self._ensure_loaded()
        idx = self._title_to_idx.get(title)
        if idx is None:
            return []

        inbound = []
        for source_idx, links in self._link_graph.items():
            if idx in links:
                inbound.append(self._titles[source_idx])

        return inbound

    # =========================================================================
    # Utility Methods
    # =========================================================================

    def validate(self) -> dict[str, bool]:
        """Run validation checks on loaded data."""
        self._ensure_loaded()
        return {
            "titles_loaded": len(self._titles) > 0,
            "title_to_idx_loaded": len(self._title_to_idx) > 0,
            "link_graph_loaded": len(self._link_graph) > 0,
            "embeddings_loaded": self._embeddings is not None,
            "faiss_index_built": self._faiss_index is not None,
            "faiss_index_count_correct": self._faiss_index.ntotal == len(self._embeddings),
            "counts_match": len(self._titles) == len(self._title_to_idx),
            "embedding_shape_correct": self._embeddings.shape[1] == EMBEDDING_DIM,
        }

    def stats(self) -> dict:
        """Get statistics about the loaded data."""
        self._ensure_loaded()
        return {
            "total_articles": len(self._titles),
            "traversable_articles": len(self._link_graph),
            "embedding_dim": self._embeddings.shape[1],
            "embedding_dtype": str(self._embeddings.dtype),
            "faiss_vectors": self._faiss_index.ntotal,
        }


# Module-level singleton instance
wiki_data = WikiData()