"""Query Expansion via Pseudo-Relevance Feedback.

Implements three query expansion strategies for improving recall:

    - RocchioExpander:      Rocchio algorithm using TF-IDF weighted vectors.
                            Expands queries based on top-ranked feedback documents.
    - SynonymExpander:      Dictionary-based synonym expansion.
    - CooccurrenceExpander: Statistical co-occurrence based term expansion.

All implementations are zero-dependency (stdlib only).
"""

from __future__ import annotations

import math
from collections import defaultdict
from typing import Dict, List, Optional, Set, Tuple


class RocchioExpander:
    """Rocchio algorithm for pseudo-relevance feedback query expansion.

    Builds TF-IDF vectors for the original query and feedback documents,
    then computes an expanded query vector:

        q' = alpha * q_orig + beta * avg(positive_docs) - gamma * avg(negative_docs)

    The top-N terms by weight from q' (excluding original query terms if
    they already appear) are returned as expansion candidates.

    Args:
        alpha:            Weight for the original query vector.
        beta:             Weight for positive feedback document centroid.
        gamma:            Weight for negative feedback document centroid.
        num_expand_terms: Maximum number of new expansion terms to return.
        num_feedback_docs: Maximum feedback documents to consider (applied
                           as a slice to the inputs).
    """

    def __init__(
        self,
        alpha: float = 1.0,
        beta: float = 0.75,
        gamma: float = 0.15,
        num_expand_terms: int = 10,
        num_feedback_docs: int = 3,
    ) -> None:
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.num_expand_terms = num_expand_terms
        self.num_feedback_docs = num_feedback_docs

    def expand(
        self,
        query_terms: List[str],
        feedback_doc_tokens: List[List[str]],
        negative_doc_tokens: Optional[List[List[str]]] = None,
    ) -> List[Tuple[str, float]]:
        """Expand query terms using the Rocchio algorithm.

        Args:
            query_terms:         Original query terms.
            feedback_doc_tokens: Token lists from top-ranked (positive) documents.
            negative_doc_tokens: Token lists from non-relevant (negative) documents.

        Returns:
            List of (term, weight) tuples sorted by descending weight.
            Includes both original query terms (re-weighted) and new
            expansion terms, up to ``num_expand_terms`` new terms beyond
            the originals.
        """
        if not query_terms:
            return []

        positive_docs = feedback_doc_tokens[: self.num_feedback_docs]
        negative_docs: List[List[str]] = []
        if negative_doc_tokens is not None:
            negative_docs = negative_doc_tokens[: self.num_feedback_docs]

        # Collect all documents for IDF computation
        all_docs = list(positive_docs) + list(negative_docs)
        if not all_docs:
            # No feedback docs — return original terms with alpha weight
            return [(t, self.alpha) for t in query_terms]

        # Build document frequency counts across the feedback corpus
        num_docs = len(all_docs)
        doc_freq: Dict[str, int] = defaultdict(int)
        for doc in all_docs:
            seen: Set[str] = set()
            for token in doc:
                if token not in seen:
                    doc_freq[token] += 1
                    seen.add(token)

        # Also count the query as a "document" for IDF
        query_set = set(query_terms)
        for t in query_set:
            doc_freq[t] = doc_freq.get(t, 0) + 1
        total_docs = num_docs + 1  # +1 for the query pseudo-document

        def _idf(term: str) -> float:
            df = doc_freq.get(term, 0)
            return math.log((total_docs + 1) / (df + 1))

        def _tfidf_vector(tokens: List[str]) -> Dict[str, float]:
            tf: Dict[str, int] = defaultdict(int)
            for t in tokens:
                tf[t] += 1
            vec: Dict[str, float] = {}
            for term, count in tf.items():
                vec[term] = (1 + math.log(count)) * _idf(term)
            return vec

        # Original query vector
        query_vec = _tfidf_vector(query_terms)

        # Positive centroid
        pos_centroid: Dict[str, float] = defaultdict(float)
        if positive_docs:
            for doc in positive_docs:
                vec = _tfidf_vector(doc)
                for term, weight in vec.items():
                    pos_centroid[term] += weight
            n_pos = len(positive_docs)
            for term in pos_centroid:
                pos_centroid[term] /= n_pos

        # Negative centroid
        neg_centroid: Dict[str, float] = defaultdict(float)
        if negative_docs:
            for doc in negative_docs:
                vec = _tfidf_vector(doc)
                for term, weight in vec.items():
                    neg_centroid[term] += weight
            n_neg = len(negative_docs)
            for term in neg_centroid:
                neg_centroid[term] /= n_neg

        # Rocchio combination: q' = alpha * q + beta * pos_centroid - gamma * neg_centroid
        expanded: Dict[str, float] = defaultdict(float)

        for term, w in query_vec.items():
            expanded[term] += self.alpha * w
        for term, w in pos_centroid.items():
            expanded[term] += self.beta * w
        for term, w in neg_centroid.items():
            expanded[term] -= self.gamma * w

        # Drop terms with non-positive weight
        expanded = {t: w for t, w in expanded.items() if w > 0}

        # Separate original terms and new expansion terms
        original = [(t, expanded[t]) for t in query_terms if t in expanded]
        new_terms = [
            (t, w)
            for t, w in expanded.items()
            if t not in query_set
        ]

        # Sort new terms by weight descending, take top N
        new_terms.sort(key=lambda x: -x[1])
        new_terms = new_terms[: self.num_expand_terms]

        # Return originals first, then expansion terms
        result = original + new_terms
        result.sort(key=lambda x: -x[1])
        return result


class SynonymExpander:
    """Dictionary-based synonym expansion.

    Maintains a mapping from terms to their synonym sets.  Supports both
    bidirectional synonym groups and directional synonym pairs.
    """

    def __init__(self) -> None:
        self._synonyms: Dict[str, Set[str]] = defaultdict(set)

    def add_synonym_group(self, words: List[str]) -> None:
        """Register a group of mutual synonyms.

        Every word in the group becomes a synonym of every other word.

        Args:
            words: List of words that are all synonyms of each other.
        """
        word_set = set(words)
        for word in words:
            self._synonyms[word].update(word_set - {word})

    def add_synonym(self, word: str, synonym: str) -> None:
        """Register a directional synonym.

        ``synonym`` is added as a synonym of ``word``, but not the reverse.

        Args:
            word:    The source term.
            synonym: The synonym to associate with ``word``.
        """
        self._synonyms[word].add(synonym)

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

        Args:
            term: The term to expand.

        Returns:
            Set of synonym strings (does not include the term itself).
        """
        return set(self._synonyms.get(term, set()))

    def expand_query(self, terms: List[str]) -> List[str]:
        """Expand all query terms with their synonyms.

        Returns the original terms followed by any expansion terms
        (deduplicated, order-preserving).

        Args:
            terms: Original query terms.

        Returns:
            Expanded list of terms.
        """
        seen: Set[str] = set()
        result: List[str] = []

        for t in terms:
            if t not in seen:
                seen.add(t)
                result.append(t)
            for syn in sorted(self._synonyms.get(t, set())):
                if syn not in seen:
                    seen.add(syn)
                    result.append(syn)

        return result

    def load_from_dict(self, synonym_dict: Dict[str, List[str]]) -> None:
        """Bulk load synonyms from a dictionary.

        Each key maps to a list of its synonyms (directional).

        Args:
            synonym_dict: Mapping of {word: [synonym, ...]}.
        """
        for word, syns in synonym_dict.items():
            for syn in syns:
                self._synonyms[word].add(syn)


class CooccurrenceExpander:
    """Statistical co-occurrence based term expansion.

    Builds a term co-occurrence matrix from documents, then expands
    query terms by finding other terms that frequently appear in the
    same documents.  Co-occurrence strength is measured by Pointwise
    Mutual Information (PMI).
    """

    def __init__(self) -> None:
        self._cooccurrence: Dict[str, Dict[str, int]] = defaultdict(
            lambda: defaultdict(int)
        )
        self._term_doc_count: Dict[str, int] = defaultdict(int)
        self._num_docs: int = 0

    def add_document(self, tokens: List[str]) -> None:
        """Record term co-occurrences from a document.

        All unique term pairs in the document contribute one
        co-occurrence count each.

        Args:
            tokens: List of tokens in the document.
        """
        self._num_docs += 1
        unique = list(set(tokens))

        for term in unique:
            self._term_doc_count[term] += 1

        for i in range(len(unique)):
            for j in range(i + 1, len(unique)):
                t1, t2 = unique[i], unique[j]
                self._cooccurrence[t1][t2] += 1
                self._cooccurrence[t2][t1] += 1

    def expand(
        self, term: str, top_k: int = 5
    ) -> List[Tuple[str, float]]:
        """Find terms that frequently co-occur with the given term.

        Ranks co-occurring terms by Pointwise Mutual Information (PMI):

            PMI(x, y) = log2(P(x,y) / (P(x) * P(y)))

        where probabilities are estimated from document frequencies.

        Args:
            term:  The term to expand.
            top_k: Maximum number of co-occurring terms to return.

        Returns:
            List of (term, pmi_score) tuples sorted by descending PMI.
        """
        if term not in self._cooccurrence or self._num_docs == 0:
            return []

        n = self._num_docs
        p_term = self._term_doc_count.get(term, 0) / n

        if p_term == 0:
            return []

        scored: List[Tuple[str, float]] = []
        for other, cocount in self._cooccurrence[term].items():
            p_other = self._term_doc_count.get(other, 0) / n
            if p_other == 0:
                continue
            p_joint = cocount / n
            pmi = math.log2(p_joint / (p_term * p_other))
            if pmi > 0:
                scored.append((other, pmi))

        scored.sort(key=lambda x: -x[1])
        return scored[:top_k]
