"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.BM25Index = void 0;
const text_1 = require("./text");
class BM25Index {
    constructor(k1 = 1.2, b = 0.75) {
        this.k1 = k1;
        this.b = b;
        this.docs = new Map();
        this.docLen = new Map();
        this.avgDocLen = 0;
        // term -> postings
        this.inverted = new Map();
        // term -> document frequency
        this.df = new Map();
    }
    addDocuments(chunks) {
        var _a, _b;
        for (const d of chunks) {
            this.docs.set(d.id, d);
            const tokens = (0, text_1.tokenize)(d.text);
            this.docLen.set(d.id, tokens.length);
            const tfMap = new Map();
            for (const t of tokens)
                tfMap.set(t, ((_a = tfMap.get(t)) !== null && _a !== void 0 ? _a : 0) + 1);
            for (const [term, tf] of tfMap) {
                const list = (_b = this.inverted.get(term)) !== null && _b !== void 0 ? _b : [];
                list.push({ docId: d.id, tf });
                this.inverted.set(term, list);
            }
        }
        // compute df and avgDocLen
        let totalLen = 0;
        for (const [docId, len] of this.docLen)
            totalLen += len;
        this.avgDocLen = this.docs.size ? totalLen / this.docs.size : 0;
        this.df.clear();
        for (const [term, postings] of this.inverted) {
            this.df.set(term, postings.length);
        }
    }
    search(query, k = 5) {
        var _a, _b, _c;
        const q = (0, text_1.tokenize)(query);
        if (q.length === 0)
            return [];
        const N = this.docs.size;
        if (N === 0)
            return [];
        const scores = new Map();
        for (const term of q) {
            const postings = this.inverted.get(term);
            if (!postings)
                continue;
            const df = (_a = this.df.get(term)) !== null && _a !== void 0 ? _a : 0;
            // standard BM25 idf
            const idf = Math.log(1 + (N - df + 0.5) / (df + 0.5));
            for (const { docId, tf } of postings) {
                const dl = (_b = this.docLen.get(docId)) !== null && _b !== void 0 ? _b : 0;
                const denom = tf + this.k1 * (1 - this.b + this.b * (dl / (this.avgDocLen || 1)));
                const score = idf * ((tf * (this.k1 + 1)) / (denom || 1));
                scores.set(docId, ((_c = scores.get(docId)) !== null && _c !== void 0 ? _c : 0) + score);
            }
        }
        return [...scores.entries()]
            .sort((a, b) => b[1] - a[1])
            .slice(0, k)
            .map(([docId, score]) => {
            const doc = this.docs.get(docId);
            return Object.assign(Object.assign({}, doc), { score });
        });
    }
}
exports.BM25Index = BM25Index;
