# retoor <retoor@molodetz.nl>
import argparse
import math
import re

K1 = 1.5
B = 0.75


def _sentences(text):
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [part for part in parts if part]


def _terms(text):
    return [word for word in re.findall(r"\w+", text.lower()) if len(word) > 2]


def select(text, query, chars=1500, window=3):
    sentences = _sentences(text)
    if not sentences:
        return ""
    query_terms = set(_terms(query))
    if not query_terms:
        return text[:chars]
    doc_freq = {}
    tokenized = [_terms(sentence) for sentence in sentences]
    for tokens in tokenized:
        for term in set(tokens) & query_terms:
            doc_freq[term] = doc_freq.get(term, 0) + 1
    count = len(sentences)
    avg_len = sum(len(tokens) for tokens in tokenized) / max(count, 1)
    idf = {term: math.log(1 + (count - freq + 0.5) / (freq + 0.5)) for term, freq in doc_freq.items()}

    def score(tokens):
        total = 0.0
        length = len(tokens)
        for term in query_terms:
            freq = tokens.count(term)
            if freq == 0:
                continue
            denom = freq + K1 * (1 - B + B * length / max(avg_len, 1))
            total += idf.get(term, 0.0) * freq * (K1 + 1) / denom
        return total

    best, best_score = (0, min(window, count)), -1.0
    for start in range(count):
        end = min(start + window, count)
        merged = [token for tokens in tokenized[start:end] for token in tokens]
        current = score(merged)
        if current > best_score:
            best, best_score = (start, end), current
    if best_score <= 0:
        return text[:chars]
    return " ".join(sentences[best[0] : best[1]])[:chars]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Select the best BM25 excerpt window for a query.")
    parser.add_argument("query", help="Research question the excerpt must answer")
    parser.add_argument("textfile", help="File holding the extracted source text")
    parser.add_argument("--chars", type=int, default=1500, help="Maximum excerpt characters")
    parser.add_argument("--window", type=int, default=3, help="Consecutive sentences per window")
    args = parser.parse_args()
    with open(args.textfile, encoding="utf-8") as handle:
        print(select(handle.read(), args.query, chars=args.chars, window=args.window))