"""Block-Max WAND (BMWAND) for efficient top-K retrieval.

An aggressive pruning algorithm that maintains per-block upper-bound scores
within each posting list, enabling entire blocks of documents to be skipped
during query evaluation.  Used by Lucene and Tantivy for production-grade
search engines.

Reference:
    Ding & Suel, "Faster Top-k Document Retrieval Using Block-Max Indexes",
    SIGIR 2011.

Algorithm overview:
    1. Sort posting lists by current doc_id.
    2. Find the pivot term where cumulative block-max scores >= threshold.
    3. If the pivot doc_id appears in all lists up to pivot, score it fully.
    4. Otherwise, advance the shortest list past the pivot using block-level
       skipping.
    5. Update the threshold from the minimum of the top-K heap.

Zero external dependencies.
"""

from __future__ import annotations

import heapq
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

BLOCK_SIZE: int = 128


# ---------------------------------------------------------------------------
# Block and BlockPostingList
# ---------------------------------------------------------------------------

@dataclass
class Block:
    """A contiguous block of postings within a posting list.

    Attributes:
        doc_ids: Sorted document IDs in this block.
        scores:  Per-document scores (parallel to doc_ids).
        max_score: Maximum score across all entries in this block.
    """

    doc_ids: List[int]
    scores: List[float]
    max_score: float


class BlockPostingList:
    """A posting list divided into fixed-size blocks with per-block max scores.

    Supports block-level skipping: ``advance_to_block`` jumps directly to the
    block containing a given doc_id, avoiding linear scans over irrelevant
    postings.

    Args:
        blocks: Pre-built list of :class:`Block` objects.  Blocks must be
            ordered by ascending doc_id with no overlap.
    """

    def __init__(self, blocks: List[Block]) -> None:
        self._blocks = blocks
        self._pos: int = 0  # current block index
        # Pre-compute global max score across all blocks
        self._max_score: float = (
            max(b.max_score for b in blocks) if blocks else 0.0
        )

    def max_score(self) -> float:
        """Maximum score across all blocks in this posting list."""
        return self._max_score

    def advance_to_block(self, min_doc_id: int) -> Optional[Block]:
        """Skip to the first block whose last doc_id >= *min_doc_id*.

        Returns the block, or ``None`` if no such block exists.
        """
        while self._pos < len(self._blocks):
            block = self._blocks[self._pos]
            # If the block's last doc_id is >= min_doc_id, this block
            # might contain it (or a later doc).
            if block.doc_ids[-1] >= min_doc_id:
                return block
            self._pos += 1
        return None

    def current_block_max(self) -> float:
        """Max score of the current block, or 0.0 if exhausted."""
        if self._pos < len(self._blocks):
            return self._blocks[self._pos].max_score
        return 0.0

    def reset(self) -> None:
        """Reset cursor to the first block."""
        self._pos = 0

    @property
    def exhausted(self) -> bool:
        """True if all blocks have been consumed."""
        return self._pos >= len(self._blocks)

    @property
    def current_block(self) -> Optional[Block]:
        """The block at the current cursor position, or None."""
        if self._pos < len(self._blocks):
            return self._blocks[self._pos]
        return None


# ---------------------------------------------------------------------------
# BlockMaxWAND evaluator
# ---------------------------------------------------------------------------

class BlockMaxWAND:
    """Block-Max WAND top-K evaluator.

    Args:
        posting_lists: Mapping from term to its :class:`BlockPostingList`.
        top_k: Number of results to return.
    """

    def __init__(
        self,
        posting_lists: Dict[str, BlockPostingList],
        top_k: int = 10,
    ) -> None:
        self._lists = posting_lists
        self._top_k = top_k

    def evaluate(self) -> List[Tuple[int, float]]:
        """Run the Block-Max WAND algorithm.

        Returns:
            Top-K ``(doc_id, score)`` pairs sorted by score descending.
        """
        if not self._lists:
            return []

        top_k = self._top_k

        # Reset all posting lists to the start
        for pl in self._lists.values():
            pl.reset()

        # Build per-term cursor state:
        #   [term, posting_list, current_doc_id]
        # We maintain a list of cursors sorted by current_doc_id.
        cursors: List[List] = []  # [term, pl, cur_doc_id]
        for term, pl in self._lists.items():
            block = pl.advance_to_block(0)
            if block is not None:
                cursors.append([term, pl, block.doc_ids[0]])

        if not cursors:
            return []

        # Min-heap of (score, doc_id) — we negate score for max-heap via min-heap
        heap: List[Tuple[float, int]] = []
        threshold: float = 0.0

        max_iterations = 0
        # Safety bound: sum of all postings across all lists
        for pl in self._lists.values():
            for b in pl._blocks:
                max_iterations += len(b.doc_ids)
        max_iterations = max(max_iterations * 2, 10000)

        iteration = 0
        while cursors and iteration < max_iterations:
            iteration += 1

            # 1. Sort cursors by current doc_id
            cursors.sort(key=lambda c: c[2])

            # 2. Find pivot: first position where cumulative block-max >= threshold
            cumulative = 0.0
            pivot_idx = -1
            for i, (term, pl, _doc_id) in enumerate(cursors):
                cumulative += pl.current_block_max()
                if cumulative >= threshold:
                    pivot_idx = i
                    break

            if pivot_idx == -1:
                # No combination of remaining terms can reach threshold
                break

            pivot_doc_id = cursors[pivot_idx][2]

            # 3. Check if all cursors up to pivot point to the same doc_id
            all_at_pivot = all(
                cursors[i][2] == pivot_doc_id for i in range(pivot_idx + 1)
            )

            if all_at_pivot:
                # Score this document fully across all terms
                score = self._score_document(pivot_doc_id)
                if score > threshold or len(heap) < top_k:
                    if len(heap) < top_k:
                        heapq.heappush(heap, (score, pivot_doc_id))
                        if len(heap) == top_k:
                            threshold = heap[0][0]
                    elif score > heap[0][0]:
                        heapq.heapreplace(heap, (score, pivot_doc_id))
                        threshold = heap[0][0]

                # Advance all cursors past pivot_doc_id
                next_doc = pivot_doc_id + 1
                alive = []
                for cursor in cursors:
                    _term, pl, _cur = cursor
                    block = pl.advance_to_block(next_doc)
                    if block is not None:
                        # Find the first doc_id >= next_doc in this block
                        new_doc = self._find_ge(block.doc_ids, next_doc)
                        if new_doc is not None:
                            cursor[2] = new_doc
                            alive.append(cursor)
                        else:
                            # Try next block
                            pl._pos += 1
                            block = pl.advance_to_block(next_doc)
                            if block is not None:
                                cursor[2] = block.doc_ids[0]
                                alive.append(cursor)
                    # else: this list is exhausted
                cursors = alive
            else:
                # 4. Advance the cursor with the smallest doc_id to pivot_doc_id
                # (using block-level skipping)
                cursor = cursors[0]
                _term, pl, _cur = cursor
                block = pl.advance_to_block(pivot_doc_id)
                if block is not None:
                    new_doc = self._find_ge(block.doc_ids, pivot_doc_id)
                    if new_doc is not None:
                        cursor[2] = new_doc
                    else:
                        pl._pos += 1
                        block = pl.advance_to_block(pivot_doc_id)
                        if block is not None:
                            cursor[2] = block.doc_ids[0]
                        else:
                            cursors.pop(0)
                else:
                    cursors.pop(0)

        # Return results sorted by score descending
        results = [(doc_id, score) for score, doc_id in heap]
        results.sort(key=lambda x: x[1], reverse=True)
        return results

    def _score_document(self, doc_id: int) -> float:
        """Compute the exact score for a document across all posting lists."""
        total = 0.0
        for _term, pl in self._lists.items():
            # Search all blocks for this doc_id
            for block in pl._blocks:
                if not block.doc_ids:
                    continue
                if block.doc_ids[0] > doc_id:
                    break
                if block.doc_ids[-1] < doc_id:
                    continue
                idx = self._bisect_find(block.doc_ids, doc_id)
                if idx is not None:
                    total += block.scores[idx]
                    break
        return total

    @staticmethod
    def _find_ge(sorted_list: List[int], target: int) -> Optional[int]:
        """Find the first value >= target in a sorted list, or None."""
        lo, hi = 0, len(sorted_list)
        while lo < hi:
            mid = (lo + hi) >> 1
            if sorted_list[mid] < target:
                lo = mid + 1
            else:
                hi = mid
        if lo < len(sorted_list):
            return sorted_list[lo]
        return None

    @staticmethod
    def _bisect_find(sorted_list: List[int], target: int) -> Optional[int]:
        """Binary search for target in sorted_list. Returns index or None."""
        lo, hi = 0, len(sorted_list)
        while lo < hi:
            mid = (lo + hi) >> 1
            if sorted_list[mid] < target:
                lo = mid + 1
            else:
                hi = mid
        if lo < len(sorted_list) and sorted_list[lo] == target:
            return lo
        return None


# ---------------------------------------------------------------------------
# Builder helper
# ---------------------------------------------------------------------------

def build_block_posting_list(
    doc_ids: List[int],
    scores: List[float],
    block_size: int = BLOCK_SIZE,
) -> BlockPostingList:
    """Build a :class:`BlockPostingList` from raw posting data.

    Args:
        doc_ids: Sorted document IDs.
        scores:  Parallel list of per-document scores.
        block_size: Number of postings per block (default 128).

    Returns:
        A new :class:`BlockPostingList` ready for query evaluation.

    Raises:
        ValueError: If doc_ids and scores have different lengths.
    """
    if len(doc_ids) != len(scores):
        raise ValueError(
            f"doc_ids length ({len(doc_ids)}) != scores length ({len(scores)})"
        )

    blocks: List[Block] = []
    for start in range(0, len(doc_ids), block_size):
        end = min(start + block_size, len(doc_ids))
        block_doc_ids = doc_ids[start:end]
        block_scores = scores[start:end]
        blocks.append(Block(
            doc_ids=block_doc_ids,
            scores=block_scores,
            max_score=max(block_scores) if block_scores else 0.0,
        ))

    return BlockPostingList(blocks)
