"""Busca por palavras-chave (BM25) e fusão de rankings (RRF).

A busca vetorial falha quando a pergunta compartilha o vocabulário do versículo
mas não a sua forma de frase — "o que Jesus disse sobre amar os inimigos?" não
recuperava Mateus 5:44 com nenhum dos três modelos medidos. O BM25 cobre esse
caso pelo lado lexical.

As duas escalas são incomparáveis (cosseno em 0..1, BM25 sem limite superior),
então a fusão usa **Reciprocal Rank Fusion**: só a posição em cada ranking
importa, não a pontuação absoluta.
"""

from __future__ import annotations

import math
import re
import unicodedata
from collections import defaultdict
from typing import Hashable, Iterable, Sequence

# Palavras que aparecem em quase todo versículo: não discriminam nada e só
# consomem memória do índice invertido.
STOPWORDS: frozenset[str] = frozenset(
    """
    a ao aos aquela aquelas aquele aqueles aquilo as até com como contra da das de
    dela delas dele deles depois do dos e ela elas ele eles em entre era eram essa
    essas esse esses esta estas este estes eu foi foram há isso isto ja lhe lhes
    mais mas me mesmo meu meus muito na nao nas nem no nos nossa nossas nosso
    nossos num numa o os ou para pela pelas pelo pelos por qual quando que quem se
    sem ser sob sobre sua suas seu seus te tem tinha tu tua tuas um uma você vocês
    vos
    """.split()
)

_NON_WORD_RE = re.compile(r"[^a-z0-9]+")
_DIGITS_RE = re.compile(r"^\d+$")

# Sufixos em etapas, cada uma aplicada no máximo uma vez (da mesma forma que o
# RSLP). Dentro da etapa, o primeiro sufixo que casa vence — por isso vão do mais
# longo ao mais curto. `min` é o comprimento mínimo do radical resultante: sem
# ele, "deus" viraria "d" e "mar" viraria "m".
_STEM_STEPS: tuple[tuple[tuple[str, str, int], ...], ...] = (
    # 1. plural
    (("ões", "ão", 3), ("ães", "ão", 3), ("ais", "al", 4), ("éis", "el", 4),
     ("óis", "ol", 4), ("ns", "m", 3), ("s", "", 3)),
    # 2. advérbio
    (("mente", "", 4),),
    # 3. superlativo, diminutivo
    (("íssima", "", 4), ("íssimo", "", 4), ("zinha", "", 4), ("zinho", "", 4),
     ("inha", "", 4), ("inho", "", 4)),
    # 4. sufixos nominais e adjetivais
    (("adora", "", 4), ("ista", "", 4), ("ismo", "", 4), ("ável", "", 4),
     ("ível", "", 4), ("ença", "", 4), ("ância", "", 4), ("idade", "", 4),
     ("ura", "", 4), ("osa", "", 4), ("oso", "", 4), ("iva", "", 4),
     ("ivo", "", 4), ("ante", "", 4), ("ador", "", 4), ("eira", "", 4),
     ("eiro", "", 4), ("eza", "", 4), ("al", "", 4)),
    # 5. verbos
    (("aríamos", "", 2), ("eríamos", "", 2), ("iríamos", "", 2),
     ("ando", "", 2), ("endo", "", 2), ("indo", "", 2),
     ("aram", "", 2), ("eram", "", 2), ("iram", "", 2),
     ("arão", "", 2), ("erão", "", 2), ("irão", "", 2),
     ("aria", "", 2), ("eria", "", 2), ("iria", "", 2),
     ("asse", "", 2), ("esse", "", 2), ("isse", "", 2),
     ("aste", "", 2), ("este", "", 2), ("iste", "", 2),
     ("arei", "", 2), ("erei", "", 2), ("irei", "", 2),
     ("amos", "", 2), ("emos", "", 2), ("imos", "", 2),
     ("avam", "", 2), ("ado", "", 2), ("ido", "", 2), ("ada", "", 2),
     ("ida", "", 2), ("ara", "", 2), ("era", "", 2), ("ira", "", 2),
     ("ava", "", 2), ("iam", "", 2),
     ("ar", "", 2), ("er", "", 2), ("ir", "", 2),
     ("ou", "", 2), ("ei", "", 2), ("eu", "", 2), ("iu", "", 2),
     ("am", "", 2), ("em", "", 2), ("as", "", 2), ("es", "", 2),
     ("ia", "", 2), ("is", "", 2), ("ai", "", 2), ("ais", "", 2),
     ("eis", "", 2), ("a", "", 2), ("e", "", 2), ("o", "", 2)),
)


def stem(word: str) -> str:
    """Reduz uma palavra já normalizada ao seu radical.

    O português é muito flexionado: sem isto, "amar" na pergunta nunca casa com
    "amai" ou "amarás" no versículo, e o BM25 passa a pontuar por palavras
    incidentais. Medido no corpus real, isso deixava Mateus 5:44 em 434º lugar
    para a pergunta "o que Jesus disse sobre amar os inimigos?".
    """
    for step in _STEM_STEPS:
        for suffix, replacement, minimum in step:
            if not word.endswith(suffix):
                continue
            candidate = word[: len(word) - len(suffix)] + replacement
            if len(candidate) >= minimum:
                word = candidate
            break
    return word


def normalize_token(word: str) -> str:
    """Minúsculas e sem acentos, para que "coração" e "coracao" coincidam."""
    decomposed = unicodedata.normalize("NFKD", word)
    return "".join(ch for ch in decomposed if not unicodedata.combining(ch)).lower()


def tokenize(text: str) -> list[str]:
    """Separa o texto em termos indexáveis, normalizados e reduzidos ao radical."""
    tokens: list[str] = []
    for raw in _NON_WORD_RE.split(normalize_token(text or "")):
        if not raw or _DIGITS_RE.match(raw):
            continue
        if raw in STOPWORDS:
            continue
        tokens.append(stem(raw))
    return tokens


class BM25Index:
    """Índice BM25 (Okapi) sobre documentos de texto.

    Mantido em memória: 31 mil versículos cabem em poucos MB e a reconstrução a
    partir do banco vetorial leva cerca de um segundo.
    """

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self._doc_freqs: dict[str, int] = {}
        self._postings: dict[str, list[tuple[int, int]]] = {}
        self._doc_lengths: list[int] = []
        self._idf: dict[str, float] = {}
        self._avgdl: float = 0.0

    def __len__(self) -> int:
        return len(self._doc_lengths)

    def build(self, documents: Sequence[str]) -> None:
        """(Re)constrói o índice, descartando os documentos anteriores."""
        postings: dict[str, list[tuple[int, int]]] = defaultdict(list)
        doc_lengths: list[int] = []
        doc_freqs: dict[str, int] = {}

        for index, document in enumerate(documents):
            tokens = tokenize(document)
            doc_lengths.append(len(tokens))
            counts: dict[str, int] = defaultdict(int)
            for token in tokens:
                counts[token] += 1
            for token, count in counts.items():
                postings[token].append((index, count))
                doc_freqs[token] = doc_freqs.get(token, 0) + 1

        total = len(doc_lengths)
        self._postings = dict(postings)
        self._doc_lengths = doc_lengths
        self._doc_freqs = doc_freqs
        self._avgdl = (sum(doc_lengths) / total) if total else 0.0
        # IDF de Robertson-Spärck Jones, na forma sempre positiva usada pelo
        # Lucene: evita peso negativo para termos presentes em quase todo doc.
        self._idf = {
            token: math.log(1.0 + (total - freq + 0.5) / (freq + 0.5))
            for token, freq in doc_freqs.items()
        }

    def search(self, query: str, k: int = 10) -> list[tuple[int, float]]:
        """Devolve `[(índice_do_documento, pontuação)]` em ordem decrescente."""
        if not self._doc_lengths or not self._avgdl:
            return []

        terms = dict.fromkeys(tokenize(query))
        if not terms:
            return []

        scores: dict[int, float] = defaultdict(float)
        for term in terms:
            idf = self._idf.get(term)
            if not idf:
                continue
            for doc_index, term_freq in self._postings.get(term, ()):
                length = self._doc_lengths[doc_index]
                denominator = term_freq + self.k1 * (
                    1.0 - self.b + self.b * length / self._avgdl
                )
                scores[doc_index] += idf * (term_freq * (self.k1 + 1.0)) / denominator

        ranked = sorted(scores.items(), key=lambda item: (-item[1], item[0]))
        return ranked[: max(0, k)]


def reciprocal_rank_fusion(
    rankings: Iterable[Sequence[Hashable]], limit: int, k: int = 60
) -> list[tuple[Hashable, float]]:
    """Combina rankings por posição: `score = Σ 1/(k + posição)`.

    A posição é 1-based. Itens ausentes de um ranking simplesmente não somam
    ali — é isso que faz um resultado confirmado pelas duas buscas liderar.
    Empates são desempatados pela ordem de aparição, para o resultado ser
    determinístico entre execuções.
    """
    scores: dict[Hashable, float] = defaultdict(float)
    order: dict[Hashable, int] = {}

    for ranking in rankings:
        for position, item in enumerate(ranking, start=1):
            scores[item] += 1.0 / (k + position)
            order.setdefault(item, len(order))

    fused = sorted(scores.items(), key=lambda entry: (-entry[1], order[entry[0]]))
    return fused[: max(0, limit)]