"""
RBSP Locality-Sensitive Hashing (LSH).

Provides semantic similarity search WITHOUT AI models by combining:
    - N-gram shingling: decomposes text into overlapping character n-grams
    - MinHash signatures: compact sketch approximating Jaccard similarity
    - Band-based LSH index: fast approximate nearest-neighbor via hash banding

All operations are O(1) per query (amortised) after the index is built.
Pure Python, no external dependencies.
"""

from __future__ import annotations

import heapq
import random
from collections import defaultdict
from typing import (
    Dict,
    Hashable,
    Iterable,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
)

_MASK64 = 0xFFFFFFFFFFFFFFFF
_LARGE_PRIME = (1 << 61) - 1  # Mersenne prime 2^61 - 1


# ---------------------------------------------------------------------------
# N-gram shingling
# ---------------------------------------------------------------------------

def ngram_shingles(text: str, n: int = 3) -> Set[int]:
    """Extract the set of character-level n-gram shingles from *text*.
    Returns pre-hashed integer shingles for efficiency.
    """
    text = text.lower()
    if len(text) < n:
        return set()
    return {hash(text[i : i + n]) for i in range(len(text) - n + 1)}


def word_shingles(text: str, n: int = 2) -> Set[int]:
    """Extract overlapping word-level n-gram shingles.
    Returns pre-hashed integer shingles for efficiency.
    """
    words = text.lower().split()
    if len(words) < n:
        return set()
    return {hash(" ".join(words[i : i + n])) for i in range(len(words) - n + 1)}


# ---------------------------------------------------------------------------
# MinHash
# ---------------------------------------------------------------------------

def _generate_hash_params(
    num_hashes: int,
    seed: int = 42,
) -> List[Tuple[int, int]]:
    """Generate (a, b) pairs for universal hash family ``h(x) = (a*x + b) % p``.

    Uses a seeded RNG for reproducibility across runs.
    """
    rng = random.Random(seed)
    params: List[Tuple[int, int]] = []
    for _ in range(num_hashes):
        a = rng.randint(1, _LARGE_PRIME - 1)
        b = rng.randint(0, _LARGE_PRIME - 1)
        params.append((a, b))
    return params


class MinHasher:
    """Compute MinHash signatures for sets of shingles.

    A MinHash signature is a compact fixed-size sketch from which the
    Jaccard similarity of two sets can be estimated with bounded error.

    Args:
        num_hashes: Length of the MinHash signature vector (default 128).
                    More hashes = more accurate similarity estimates.
        seed:       RNG seed for reproducible hash function families.
    """

    def __init__(self, num_hashes: int = 128, seed: int = 42) -> None:
        self.num_hashes = num_hashes
        self._params = _generate_hash_params(num_hashes, seed)
        self._a_vals = [a for a, b in self._params]
        self._b_vals = [b for a, b in self._params]

    def signature(self, shingles: Iterable[Hashable]) -> List[int]:
        """Compute the MinHash signature for a set of shingles.

        Args:
            shingles: An iterable of hashable shingle values.

        Returns:
            A list of *num_hashes* integers representing the MinHash sketch.
        """
        sig = [_LARGE_PRIME] * self.num_hashes
        a_vals = self._a_vals
        b_vals = self._b_vals

        for shingle in shingles:
            h_val = hash(shingle) & _MASK64
            for i in range(self.num_hashes):
                mh = ((a_vals[i] * h_val + b_vals[i]) % _LARGE_PRIME) & _MASK64
                if mh < sig[i]:
                    sig[i] = mh

        return sig

    @staticmethod
    def estimate_jaccard(sig_a: Sequence[int], sig_b: Sequence[int]) -> float:
        """Estimate the Jaccard similarity from two MinHash signatures.

        Args:
            sig_a: MinHash signature of set A.
            sig_b: MinHash signature of set B.

        Returns:
            Estimated Jaccard index in [0.0, 1.0].

        Raises:
            ValueError: If signature lengths differ.
        """
        if len(sig_a) != len(sig_b):
            raise ValueError("Signature lengths must match")
        if not sig_a:
            return 0.0
        matches = sum(1 for a, b in zip(sig_a, sig_b) if a == b)
        return matches / len(sig_a)


# ---------------------------------------------------------------------------
# Band-based LSH Index
# ---------------------------------------------------------------------------

class LSHIndex:
    """Locality-Sensitive Hashing index for approximate nearest-neighbor search.

    Uses the banding technique: each MinHash signature is split into *bands*
    bands of *rows* rows.  Two items become candidates if they share at
    least one identical band.  This gives a threshold curve where pairs with
    Jaccard similarity above roughly ``(1/bands)^(1/rows)`` are likely to be
    found.

    Args:
        num_hashes: Total MinHash signature length (must equal bands * rows).
        bands:      Number of bands (default 16).
        rows:       Rows per band (default 8).
        seed:       RNG seed forwarded to MinHasher.
    """

    def __init__(
        self,
        num_hashes: int = 128,
        bands: int = 16,
        rows: int = 8,
        seed: int = 42,
    ) -> None:
        if num_hashes != bands * rows:
            raise ValueError(
                f"num_hashes ({num_hashes}) must equal bands * rows ({bands * rows})"
            )
        self.bands = bands
        self.rows = rows
        self.num_hashes = num_hashes
        self._hasher = MinHasher(num_hashes=num_hashes, seed=seed)

        # band_index -> { band_hash: set_of_ids }
        self._buckets: List[Dict[int, Set[Hashable]]] = [
            defaultdict(set) for _ in range(bands)
        ]
        # id -> minhash signature (for re-ranking)
        self._signatures: Dict[Hashable, List[int]] = {}

    @property
    def threshold(self) -> float:
        """Approximate Jaccard similarity threshold for candidate retrieval."""
        return (1.0 / self.bands) ** (1.0 / self.rows)

    def insert(self, item_id: Hashable, shingles: Iterable[Hashable]) -> List[int]:
        """Insert an item into the LSH index.

        Args:
            item_id:  A unique identifier for the item.
            shingles: The item's set of shingles.

        Returns:
            The computed MinHash signature.
        """
        sig = self._hasher.signature(shingles)
        self._signatures[item_id] = sig

        for band_idx in range(self.bands):
            start = band_idx * self.rows
            band = tuple(sig[start : start + self.rows])
            band_hash = hash(band)
            self._buckets[band_idx][band_hash].add(item_id)

        return sig

    def remove(self, item_id: Hashable) -> bool:
        """Remove an item from the index.

        Returns:
            True if the item was found and removed, False otherwise.
        """
        sig = self._signatures.pop(item_id, None)
        if sig is None:
            return False

        for band_idx in range(self.bands):
            start = band_idx * self.rows
            band = tuple(sig[start : start + self.rows])
            band_hash = hash(band)
            bucket = self._buckets[band_idx].get(band_hash)
            if bucket is not None:
                bucket.discard(item_id)
                if not bucket:
                    del self._buckets[band_idx][band_hash]
        return True

    def query(
        self,
        shingles: Iterable[Hashable],
        *,
        top_k: int = 10,
        min_similarity: float = 0.0,
    ) -> List[Tuple[Hashable, float]]:
        """Find approximate nearest neighbors for *shingles*.

        Args:
            shingles:       Query shingle set.
            top_k:          Maximum number of results (default 10).
            min_similarity: Minimum Jaccard estimate to include (default 0.0).

        Returns:
            List of ``(item_id, estimated_jaccard)`` tuples sorted by
            descending similarity.
        """
        query_sig = self._hasher.signature(shingles)

        # Candidate collection: union of all bucket matches across bands
        candidates: Set[Hashable] = set()
        for band_idx in range(self.bands):
            start = band_idx * self.rows
            band = tuple(query_sig[start : start + self.rows])
            band_hash = hash(band)
            bucket = self._buckets[band_idx].get(band_hash)
            if bucket:
                candidates.update(bucket)

        # Re-rank candidates by exact MinHash Jaccard estimate
        results: List[Tuple[Hashable, float]] = []
        for cid in candidates:
            sim = MinHasher.estimate_jaccard(query_sig, self._signatures[cid])
            if sim >= min_similarity:
                results.append((cid, sim))

        return heapq.nlargest(top_k, results, key=lambda x: x[1])

    def query_by_id(
        self,
        item_id: Hashable,
        *,
        top_k: int = 10,
        min_similarity: float = 0.0,
    ) -> List[Tuple[Hashable, float]]:
        """Find approximate nearest neighbors of an already-indexed item.

        Raises:
            KeyError: If *item_id* is not in the index.
        """
        sig = self._signatures.get(item_id)
        if sig is None:
            raise KeyError(f"Item {item_id!r} not found in index")

        candidates: Set[Hashable] = set()
        for band_idx in range(self.bands):
            start = band_idx * self.rows
            band = tuple(sig[start : start + self.rows])
            band_hash = hash(band)
            bucket = self._buckets[band_idx].get(band_hash)
            if bucket:
                candidates.update(bucket)

        candidates.discard(item_id)

        results: List[Tuple[Hashable, float]] = []
        for cid in candidates:
            sim = MinHasher.estimate_jaccard(sig, self._signatures[cid])
            if sim >= min_similarity:
                results.append((cid, sim))

        return heapq.nlargest(top_k, results, key=lambda x: x[1])

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

    def __contains__(self, item_id: Hashable) -> bool:
        return item_id in self._signatures
