from __future__ import annotations

import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Iterable

TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")


def tokenize(text: str) -> list[str]:
    return [match.group(0).lower() for match in TOKEN_RE.finditer(text)]


@dataclass
class LexicalDoc:
    term_freq: dict[str, float]
    norm: float


class LexicalIndex:
    def __init__(self, docs: Iterable[str]) -> None:
        texts = list(docs)
        self.doc_count = len(texts)
        self.doc_freq: Counter[str] = Counter()
        self.docs: list[LexicalDoc] = []

        tokenized = [tokenize(text) for text in texts]
        for tokens in tokenized:
            self.doc_freq.update(set(tokens))

        for tokens in tokenized:
            tf = Counter(tokens)
            weighted: dict[str, float] = {}
            for term, freq in tf.items():
                idf = self._idf(term)
                weighted[term] = float(freq) * idf
            norm = math.sqrt(sum(value * value for value in weighted.values())) or 1.0
            self.docs.append(LexicalDoc(term_freq=weighted, norm=norm))

    def _idf(self, term: str) -> float:
        return math.log((self.doc_count + 1) / (self.doc_freq.get(term, 0) + 1)) + 1.0

    def search(self, query: str, k: int) -> list[tuple[int, float]]:
        if self.doc_count == 0:
            return []
        q_tokens = tokenize(query)
        if not q_tokens:
            return []
        q_tf = Counter(q_tokens)
        q_vec: dict[str, float] = {}
        for term, freq in q_tf.items():
            q_vec[term] = float(freq) * self._idf(term)
        q_norm = math.sqrt(sum(value * value for value in q_vec.values())) or 1.0

        scores: list[tuple[int, float]] = []
        for idx, doc in enumerate(self.docs):
            score = 0.0
            for term, q_val in q_vec.items():
                if term in doc.term_freq:
                    score += q_val * doc.term_freq[term]
            score = score / (q_norm * doc.norm)
            if score > 0:
                scores.append((idx, score))

        scores.sort(key=lambda item: item[1], reverse=True)
        return scores[:k]
