import re
from collections import Counter
from typing import Dict, List

import numpy as np

__all__ = ["SimpleBM25"]


class SimpleBM25:
    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.doc_tokens: List[List[str]] = []
        self.avg_dl: float = 0
        self.df: Dict[str, int] = {}
        self.n_docs: int = 0

    def fit(self, documents: List[str]):
        self.doc_tokens = [self._tokenize(d) for d in documents]
        self.n_docs = len(self.doc_tokens)
        self.avg_dl = np.mean([len(t) for t in self.doc_tokens]) if self.doc_tokens else 1
        self.df = Counter()
        for tokens in self.doc_tokens:
            for t in set(tokens):
                self.df[t] += 1

    def _tokenize(self, text: str) -> List[str]:
        return re.findall(r'\b[a-z]+\b', str(text).lower())

    def score_candidates(self, query: str, candidate_indices: List[int]) -> np.ndarray:
        q_tokens = self._tokenize(query)
        scores = np.zeros(len(candidate_indices), dtype=np.float32)
        for i, doc_idx in enumerate(candidate_indices):
            if doc_idx >= len(self.doc_tokens):
                continue
            doc = self.doc_tokens[doc_idx]
            dl = len(doc)
            tf_doc = Counter(doc)
            s = 0.0
            for qt in q_tokens:
                if qt not in self.df:
                    continue
                tf = tf_doc.get(qt, 0)
                idf = np.log((self.n_docs - self.df[qt] + 0.5) / (self.df[qt] + 0.5) + 1)
                num = tf * (self.k1 + 1)
                den = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
                s += idf * num / den
            scores[i] = s
        return scores