"""
RBSP Term Proximity Scoring.

Boosts documents where query terms appear close together.  Uses a
decay-based model: for each pair of query terms the minimum positional
span is computed and the pair contributes ``decay^(span - 1)`` to the
overall proximity score.

Minimum-span algorithm:
    Maintain one pointer per term's position list.  At each step advance
    the pointer sitting at the overall minimum position.  Track the
    smallest window seen across all iterations.

Components:
    - ProximityScorer: Configurable scorer with decay factor.
    - SpanQuery:       Dataclass describing a multi-term span constraint.
    - span_match:      Returns positions where all terms appear within a
                       maximum span window.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List


@dataclass(frozen=True, slots=True)
class SpanQuery:
    """A span constraint requiring terms to appear within a window.

    Attributes:
        terms:    Terms that must co-occur within the span.
        max_span: Maximum positional window size (inclusive).
    """

    terms: list[str]
    max_span: int


class ProximityScorer:
    """Scores documents by how close query terms appear to each other.

    For every unique pair of query terms present in the document, the
    scorer finds their minimum positional distance and awards
    ``decay^(distance - 1)``.  The total is normalised by the number of
    term pairs so the raw score stays in a predictable range.

    Args:
        decay: Multiplicative decay per position gap (default 0.9).
               A higher value means distant terms are penalised less.
    """

    def __init__(self, decay: float = 0.9) -> None:
        if not 0.0 < decay <= 1.0:
            raise ValueError("decay must be in (0, 1]")
        self._decay = decay

    def score(
        self,
        query_terms: list[str],
        term_positions: dict[str, list[int]],
    ) -> float:
        """Compute a raw proximity score for one document.

        Args:
            query_terms:    Ordered list of query terms.
            term_positions: Mapping of term -> sorted positions in the
                            document.  Terms absent from the document
                            should be omitted (or have an empty list).

        Returns:
            Sum of ``decay^(min_span - 1)`` over each unique term pair,
            divided by the number of pairs.  Returns 0.0 when fewer than
            two terms have positions.
        """
        present = [t for t in query_terms if term_positions.get(t)]
        if len(present) < 2:
            return 0.0

        n_pairs = 0
        total = 0.0

        for i in range(len(present)):
            for j in range(i + 1, len(present)):
                pos_a = term_positions[present[i]]
                pos_b = term_positions[present[j]]
                span = _min_distance(pos_a, pos_b)
                total += self._decay ** (span - 1)
                n_pairs += 1

        return total / n_pairs

    def score_normalized(
        self,
        query_terms: list[str],
        term_positions: dict[str, list[int]],
    ) -> float:
        """Compute a proximity score normalised to [0, 1].

        Uses saturating normalisation: ``raw / (raw + 1)``.

        Args:
            query_terms:    Ordered list of query terms.
            term_positions: Mapping of term -> sorted positions in the
                            document.

        Returns:
            Normalised score in [0.0, 1.0].
        """
        raw = self.score(query_terms, term_positions)
        if raw <= 0.0:
            return 0.0
        return raw / (raw + 1.0)

    @staticmethod
    def minimum_span(positions_list: list[list[int]]) -> int:
        """Find the minimum window containing at least one position from
        every list.

        Uses a pointer-advancing technique: keep one pointer per list,
        always advance the pointer at the global minimum, and track the
        smallest ``max - min`` seen.

        Args:
            positions_list: List of sorted (ascending) position lists.
                            Every list must be non-empty.

        Returns:
            The size of the smallest window (max_pos - min_pos).

        Raises:
            ValueError: If any position list is empty.
        """
        n = len(positions_list)
        if n == 0:
            return 0
        if n == 1:
            return 0

        for pl in positions_list:
            if not pl:
                raise ValueError("All position lists must be non-empty")

        pointers = [0] * n
        best = float("inf")

        # Track the current maximum across all pointers
        current_max = max(pl[0] for pl in positions_list)

        while True:
            # Find the pointer at the minimum position
            min_val = positions_list[0][pointers[0]]
            min_idx = 0
            for k in range(1, n):
                val = positions_list[k][pointers[k]]
                if val < min_val:
                    min_val = val
                    min_idx = k

            span = current_max - min_val
            if span < best:
                best = span

            # Advance the minimum pointer
            pointers[min_idx] += 1
            if pointers[min_idx] >= len(positions_list[min_idx]):
                break

            new_val = positions_list[min_idx][pointers[min_idx]]
            if new_val > current_max:
                current_max = new_val

        return best


def span_match(
    term_positions: dict[str, list[int]],
    terms: list[str],
    max_span: int,
) -> list[int]:
    """Find document positions where all *terms* appear within *max_span*.

    Returns the starting positions (minimum position of the window) for
    every qualifying window.

    Algorithm: pointer-advance across all terms' position lists.  At each
    step compute the current window; if it fits within ``max_span`` record
    the minimum position then advance the minimum pointer to find the
    next candidate window.

    Args:
        term_positions: Mapping of term -> sorted positions.
        terms:          Terms that must all appear in the window.
        max_span:       Maximum window size (max_pos - min_pos).

    Returns:
        Sorted list of start positions for qualifying windows.
    """
    if not terms:
        return []

    positions_list: list[list[int]] = []
    for t in terms:
        pl = term_positions.get(t)
        if not pl:
            return []
        positions_list.append(pl)

    n = len(positions_list)
    if n == 1:
        # Every position of the single term qualifies
        return list(positions_list[0])

    pointers = [0] * n
    current_max = max(pl[0] for pl in positions_list)
    results: list[int] = []

    while True:
        # Find pointer at minimum position
        min_val = positions_list[0][pointers[0]]
        min_idx = 0
        for k in range(1, n):
            val = positions_list[k][pointers[k]]
            if val < min_val:
                min_val = val
                min_idx = k

        if current_max - min_val <= max_span:
            results.append(min_val)

        # Advance the minimum pointer
        pointers[min_idx] += 1
        if pointers[min_idx] >= len(positions_list[min_idx]):
            break

        new_val = positions_list[min_idx][pointers[min_idx]]
        if new_val > current_max:
            current_max = new_val

    return results


# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------

def _min_distance(sorted_a: list[int], sorted_b: list[int]) -> int:
    """Find the minimum absolute distance between two sorted lists.

    Runs in O(n + m) via a merge-style scan.

    Args:
        sorted_a: Sorted ascending positions for term A.
        sorted_b: Sorted ascending positions for term B.

    Returns:
        Minimum |a - b| across all pairs.
    """
    i = 0
    j = 0
    best = abs(sorted_a[0] - sorted_b[0])

    while i < len(sorted_a) and j < len(sorted_b):
        diff = abs(sorted_a[i] - sorted_b[j])
        if diff < best:
            best = diff
        if best == 0:
            return 0
        if sorted_a[i] < sorted_b[j]:
            i += 1
        else:
            j += 1

    return best
