"""Semantic Cluster Engine for RBSP.

Groups related signatures via LSH bucket IDs, enabling O(1) cluster retrieval
and semantic-proximity queries.  When a signature is indexed, its LSH band
hashes are used to assign it to one or more clusters.  Items that share
clusters are semantically similar.

Uses the band-based LSH approach from :mod:`rbsp.core.lsh` but organises the
results into a cluster data structure optimised for dynamic updates.

Thread-safe for concurrent reads and writes.
"""

from __future__ import annotations

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


class SemanticClusterEngine:
    """Groups items into clusters based on LSH bucket assignments.

    Each item is associated with a set of bucket IDs (typically computed by
    the LSH banding step).  Items that share bucket IDs are placed in the
    same cluster.

    Provides O(1) cluster retrieval by bucket ID and O(1) lookup of an
    item's cluster memberships.

    Usage::

        engine = SemanticClusterEngine()
        engine.add("sig_1", buckets=(3, 17, 42, 99))
        engine.add("sig_2", buckets=(3, 17, 55, 99))
        neighbours = engine.get_neighbours("sig_1")  # includes "sig_2"
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        # bucket_id -> set of item_ids in that bucket
        self._bucket_to_items: Dict[int, Set[str]] = defaultdict(set)
        # item_id -> frozenset of bucket_ids it belongs to
        self._item_buckets: Dict[str, FrozenSet[int]] = {}

    # ------------------------------------------------------------------
    # Mutation
    # ------------------------------------------------------------------

    def add(self, item_id: str, buckets: Iterable[int]) -> None:
        """Add or update an item's cluster memberships.

        If the item already exists its old bucket assignments are replaced.

        Args:
            item_id: Unique identifier (typically a signature ID).
            buckets: Iterable of integer LSH bucket IDs.
        """
        bucket_set = frozenset(buckets)
        with self._lock:
            # Remove old assignments if re-indexing
            old = self._item_buckets.get(item_id)
            if old is not None:
                for b in old:
                    self._bucket_to_items[b].discard(item_id)
                    if not self._bucket_to_items[b]:
                        del self._bucket_to_items[b]
            # Add new assignments
            self._item_buckets[item_id] = bucket_set
            for b in bucket_set:
                self._bucket_to_items[b].add(item_id)

    def remove(self, item_id: str) -> bool:
        """Remove an item from all clusters.

        Returns:
            True if the item existed and was removed.
        """
        with self._lock:
            old = self._item_buckets.pop(item_id, None)
            if old is None:
                return False
            for b in old:
                self._bucket_to_items[b].discard(item_id)
                if not self._bucket_to_items[b]:
                    del self._bucket_to_items[b]
            return True

    # ------------------------------------------------------------------
    # Lookup
    # ------------------------------------------------------------------

    def get_cluster(self, bucket_id: int) -> Set[str]:
        """O(1) retrieval of all item IDs in a given bucket/cluster.

        Args:
            bucket_id: The LSH bucket ID.

        Returns:
            A copy of the set of item IDs in this bucket.
        """
        with self._lock:
            return set(self._bucket_to_items.get(bucket_id, set()))

    def get_item_buckets(self, item_id: str) -> FrozenSet[int]:
        """Return the bucket IDs assigned to *item_id*, or empty frozenset."""
        with self._lock:
            return self._item_buckets.get(item_id, frozenset())

    def get_neighbours(self, item_id: str, *, exclude_self: bool = True) -> Set[str]:
        """Find all items that share at least one bucket with *item_id*.

        This is the primary semantic-proximity query: returns all items
        deemed similar by the LSH bucketing.

        Args:
            item_id:      The item to query neighbours for.
            exclude_self: Whether to exclude *item_id* from the result.

        Returns:
            Set of neighbouring item IDs.
        """
        with self._lock:
            buckets = self._item_buckets.get(item_id)
            if buckets is None:
                return set()
            result: Set[str] = set()
            for b in buckets:
                result.update(self._bucket_to_items.get(b, set()))
            if exclude_self:
                result.discard(item_id)
            return result

    def get_neighbours_ranked(
        self,
        item_id: str,
        *,
        top_k: int = 10,
        exclude_self: bool = True,
    ) -> List[Tuple[str, int]]:
        """Find neighbours ranked by the number of shared buckets.

        More shared buckets implies higher semantic similarity.

        Args:
            item_id:      The item to query.
            top_k:        Maximum results to return.
            exclude_self: Whether to exclude *item_id*.

        Returns:
            List of ``(neighbour_id, shared_count)`` sorted descending by
            shared bucket count.
        """
        with self._lock:
            buckets = self._item_buckets.get(item_id)
            if buckets is None:
                return []
            # Count shared buckets per neighbour
            counts: Dict[str, int] = defaultdict(int)
            for b in buckets:
                for nid in self._bucket_to_items.get(b, set()):
                    counts[nid] += 1
            if exclude_self:
                counts.pop(item_id, None)

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

    def query_by_buckets(
        self,
        buckets: Iterable[int],
        *,
        top_k: int = 10,
    ) -> List[Tuple[str, int]]:
        """Query the cluster engine directly with bucket IDs.

        Useful when you have LSH bucket IDs but no item_id in the index.

        Args:
            buckets: Iterable of LSH bucket IDs.
            top_k:   Maximum results to return.

        Returns:
            List of ``(item_id, shared_count)`` sorted descending.
        """
        bucket_list = list(buckets)
        counts: Dict[str, int] = defaultdict(int)
        with self._lock:
            for b in bucket_list:
                for nid in self._bucket_to_items.get(b, set()):
                    counts[nid] += 1
        return heapq.nlargest(top_k, counts.items(), key=lambda x: x[1])

    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------

    def item_count(self) -> int:
        """Total number of indexed items."""
        with self._lock:
            return len(self._item_buckets)

    def bucket_count(self) -> int:
        """Number of non-empty buckets."""
        with self._lock:
            return len(self._bucket_to_items)

    def cluster_size(self, bucket_id: int) -> int:
        """Number of items in a specific bucket."""
        with self._lock:
            return len(self._bucket_to_items.get(bucket_id, set()))

    def __len__(self) -> int:
        return self.item_count()

    def __contains__(self, item_id: str) -> bool:
        with self._lock:
            return item_id in self._item_buckets

    def clear(self) -> None:
        """Remove all items and clusters."""
        with self._lock:
            self._bucket_to_items.clear()
            self._item_buckets.clear()
