"""Synonym Dictionary for RBSP.

Provides query-time and index-time synonym expansion using a union-find
data structure for equivalence classes.  Each equivalence class has a
canonical form (the first term added to that group).

Supports:
    - Bidirectional and directional synonym rules
    - Equivalence groups (all terms interchangeable)
    - OR-expansion for query-time broadening
    - Canonical replacement for index-time normalization
    - Persistent storage via simple text files

Zero external dependencies.
"""

from __future__ import annotations

from typing import Dict, List, Set


# ---------------------------------------------------------------------------
# Union-Find (Disjoint Set Union) — internal helper
# ---------------------------------------------------------------------------

class UnionFind:
    """Weighted union-find with path compression.

    Tracks equivalence classes of string terms.  The canonical form
    of each group is the first term that was added to the group.
    """

    __slots__ = ("_parent", "_rank", "_canonical")

    def __init__(self) -> None:
        self._parent: Dict[str, str] = {}
        self._rank: Dict[str, int] = {}
        self._canonical: Dict[str, str] = {}

    def _ensure(self, x: str) -> None:
        """Lazily initialise a node the first time it appears."""
        if x not in self._parent:
            self._parent[x] = x
            self._rank[x] = 0
            self._canonical[x] = x

    def find(self, x: str) -> str:
        """Find root of *x* with path compression."""
        self._ensure(x)
        root = x
        while self._parent[root] != root:
            root = self._parent[root]
        # Path compression
        while self._parent[x] != root:
            nxt = self._parent[x]
            self._parent[x] = root
            x = nxt
        return root

    def union(self, x: str, y: str) -> None:
        """Merge the equivalence classes of *x* and *y*.

        The canonical form is preserved from the group that was created
        first (i.e. the group that already holds the canonical label).
        """
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return
        # Union by rank
        if self._rank[rx] < self._rank[ry]:
            rx, ry = ry, rx
        self._parent[ry] = rx
        if self._rank[rx] == self._rank[ry]:
            self._rank[rx] += 1
        # Keep the earlier canonical form
        cx, cy = self._canonical[rx], self._canonical[ry]
        # Prefer the one that was added as a canonical term first;
        # since both are already roots' canonicals, keep rx's.
        # (caller controls order — first term in add_equivalent wins)
        self._canonical[rx] = cx

    def group(self, x: str) -> Set[str]:
        """Return all members of *x*'s equivalence class."""
        self._ensure(x)
        root = self.find(x)
        return {k for k in self._parent if self.find(k) == root}

    def canonical(self, x: str) -> str:
        """Return the canonical form for *x*'s group."""
        root = self.find(x)
        return self._canonical[root]

    def __contains__(self, x: str) -> bool:
        return x in self._parent


# ---------------------------------------------------------------------------
# SynonymMap — public API
# ---------------------------------------------------------------------------

class SynonymMap:
    """Synonym dictionary with equivalence-class semantics.

    Maintains a union-find structure of equivalent terms.  Supports
    both symmetric (bidirectional) and asymmetric (directional) rules,
    query-time OR-expansion, and index-time canonical replacement.

    Example::

        sm = SynonymMap()
        sm.add_equivalent(["usa", "united states", "america"])
        sm.expand("usa")
        # {"usa", "united states", "america"}
        sm.apply_replacements(["visit", "america"])
        # ["visit", "usa"]
    """

    def __init__(self) -> None:
        self._uf = UnionFind()
        self._directional: Dict[str, Set[str]] = {}
        self._rule_count: int = 0

    # -- Rule creation ------------------------------------------------------

    def add_rule(
        self,
        from_terms: List[str],
        to_terms: List[str],
        bidirectional: bool = True,
    ) -> None:
        """Map one set of terms to another.

        Args:
            from_terms:    Source terms.
            to_terms:      Destination terms.
            bidirectional: If ``True`` (default), terms on both sides are
                           treated as fully equivalent.  If ``False``, only
                           from -> to mappings are created.
        """
        if bidirectional:
            # Merge all terms into one equivalence class
            all_terms = from_terms + to_terms
            for t in all_terms[1:]:
                self._uf.union(all_terms[0], t)
        else:
            # Directional: from-terms expand to include to-terms
            for ft in from_terms:
                self._uf._ensure(ft)
                if ft not in self._directional:
                    self._directional[ft] = set()
                for tt in to_terms:
                    self._uf._ensure(tt)
                    self._directional[ft].add(tt)
        self._rule_count += 1

    def add_equivalent(self, terms: List[str]) -> None:
        """Declare all *terms* as equivalent synonyms.

        The first term becomes the canonical form for the group.

        Args:
            terms: Two or more equivalent terms
                   (e.g. ``["usa", "united states", "america"]``).
        """
        if len(terms) < 2:
            return
        # Ensure the first term is inserted first so it becomes canonical
        self._uf._ensure(terms[0])
        for t in terms[1:]:
            self._uf.union(terms[0], t)
        # Force canonical to be the first term provided
        root = self._uf.find(terms[0])
        self._uf._canonical[root] = terms[0]
        self._rule_count += 1

    # -- Expansion ----------------------------------------------------------

    def expand(self, term: str) -> Set[str]:
        """Get all synonyms for a single term (including itself).

        Returns the equivalence class from union-find merged with any
        directional mappings.

        Args:
            term: The term to expand.

        Returns:
            Set of all synonyms (always includes *term* itself).
        """
        result: Set[str] = set()
        if term in self._uf:
            result = self._uf.group(term)
        else:
            result.add(term)
        # Add directional targets
        if term in self._directional:
            result.update(self._directional[term])
        return result

    def expand_query(self, tokens: List[str]) -> List[List[str]]:
        """Expand each token into a list of alternatives for OR-expansion.

        For each input token, returns ``[original, syn1, syn2, ...]``.
        Tokens without synonyms yield a single-element list.

        Args:
            tokens: List of query tokens.

        Returns:
            List of lists, one per input token.
        """
        result: List[List[str]] = []
        for token in tokens:
            synonyms = self.expand(token)
            # Ensure original token is first
            alternatives = [token] + sorted(synonyms - {token})
            result.append(alternatives)
        return result

    def apply_replacements(self, tokens: List[str]) -> List[str]:
        """Replace each token with its canonical form.

        Useful for index-time normalization so that all variants of a
        term map to the same posting list.

        Args:
            tokens: List of tokens.

        Returns:
            List of tokens with synonyms replaced by their canonical form.
        """
        result: List[str] = []
        for token in tokens:
            if token in self._uf:
                result.append(self._uf.canonical(token))
            else:
                result.append(token)
        return result

    # -- Persistence --------------------------------------------------------

    def load_from_file(self, path: str) -> None:
        """Load synonym rules from a text file.

        File format — one rule per line, comma-separated terms::

            usa, united states, america
            colour, color
            => big, large, huge

        Lines starting with ``=>`` are directional (first term maps to
        the rest).  Blank lines and lines starting with ``#`` are ignored.

        Args:
            path: Path to the synonym file.
        """
        with open(path, "r", encoding="utf-8") as fh:
            for raw_line in fh:
                line = raw_line.strip()
                if not line or line.startswith("#"):
                    continue
                if line.startswith("=>"):
                    # Directional rule: "=> from, to1, to2, ..."
                    parts = [t.strip() for t in line[2:].split(",") if t.strip()]
                    if len(parts) >= 2:
                        self.add_rule([parts[0]], parts[1:], bidirectional=False)
                else:
                    # Equivalence rule
                    parts = [t.strip() for t in line.split(",") if t.strip()]
                    if len(parts) >= 2:
                        self.add_equivalent(parts)

    def save_to_file(self, path: str) -> None:
        """Save synonym rules to a text file.

        Writes equivalence classes and directional rules in the format
        understood by :meth:`load_from_file`.

        Args:
            path: Destination file path.
        """
        seen_roots: Set[str] = set()
        lines: List[str] = []

        # Equivalence groups
        for term in sorted(self._uf._parent):
            root = self._uf.find(term)
            if root in seen_roots:
                continue
            seen_roots.add(root)
            members = self._uf.group(term)
            if len(members) >= 2:
                canonical = self._uf.canonical(term)
                others = sorted(members - {canonical})
                lines.append(", ".join([canonical] + others))

        # Directional rules
        for ft in sorted(self._directional):
            targets = sorted(self._directional[ft])
            if targets:
                lines.append("=> " + ", ".join([ft] + targets))

        with open(path, "w", encoding="utf-8") as fh:
            fh.write("\n".join(lines))
            if lines:
                fh.write("\n")

    # -- Dunder -------------------------------------------------------------

    def __len__(self) -> int:
        """Number of synonym rules that have been added."""
        return self._rule_count

    def __repr__(self) -> str:
        return f"SynonymMap(rules={self._rule_count})"
