data-sampler

/**
 * @file relevance.ts
 * @description Relevance-based sampling strategy.
 * Scores chunks by BM25-style term frequency against a user-provided query.
 * Returns chunks sorted by relevance score descending.
 */

import type { Chunk, SamplingConfig } from "../types";
import { sampleByLength } from "./length";

/**
 * BM25 k1 parameter — controls term frequency saturation.
 * Standard value: 1.2 to 2.0.
 */
const BM25_K1 = 1.5;

/**
 * BM25 b parameter — controls document length normalisation.
 * 0 = no length normalisation, 1 = full normalisation.
 */
const BM25_B = 0.75;

/**
 * Average chunk length in characters used as BM25 document-length normaliser.
 * When the actual average differs, BM25 adjusts term saturation accordingly.
 */
const AVG_CHUNK_LENGTH = 2000;

// ---------------------------------------------------------------------------
// Text utilities
// ---------------------------------------------------------------------------

/**
 * Tokenise text into lower-cased word tokens.
 * Words must be at least 2 characters long.
 */
function tokenise(text: string): readonly string[] {
  return text.toLowerCase().match(/\p{L}{2,}/gu) ?? [];
}

/**
 * Count term frequencies in a token list.
 */
function termFrequencies(
  tokens: readonly string[],
): ReadonlyMap<string, number> {
  const map = new Map<string, number>();
  for (const t of tokens) {
    map.set(t, (map.get(t) ?? 0) + 1);
  }
  return map;
}

/**
 * Count how many chunks contain a given term (document frequency).
 */
function documentFrequencies(
  chunks: readonly (readonly string[])[],
): ReadonlyMap<string, number> {
  const map = new Map<string, number>();
  for (const tokens of chunks) {
    const seen = new Set(tokens);
    for (const t of seen) {
      map.set(t, (map.get(t) ?? 0) + 1);
    }
  }
  return map;
}

// ---------------------------------------------------------------------------
// BM25 scoring
// ---------------------------------------------------------------------------

/**
 * Compute BM25 score for a single chunk against the query terms.
 *
 * BM25 formula:
 *   score = sum over query terms of IDF(q) * (f(q,D) * (k1 + 1)) /
 *           (f(q,D) + k1 * (1 - b + b * |D| / avgdl))
 *
 * Where:
 *   IDF(q) = ln(1 + (N - n(q) + 0.5) / (n(q) + 0.5))   — smooth IDF
 *   f(q,D) = term frequency of q in chunk D
 *   N      = total number of chunks
 *   n(q)   = number of chunks containing q
 *   |D|    = character length of chunk D
 *   avgdl  = average chunk length
 */
function bm25Score(
  chunkText: string,
  chunkLength: number,
  queryTokens: readonly string[],
  df: ReadonlyMap<string, number>,
  numChunks: number,
  avgdl: number,
): number {
  const tf = termFrequencies(tokenise(chunkText));
  let score = 0;

  for (const qt of queryTokens) {
    const fq = tf.get(qt) ?? 0;
    if (fq === 0) continue;

    const nq = df.get(qt) ?? 0;
    if (nq === 0) continue;

    // Smooth IDF
    const idf = Math.log(1 + (numChunks - nq + 0.5) / (nq + 0.5));

    // Length normalised term frequency
    const numerator = fq * (BM25_K1 + 1);
    const denom = fq + BM25_K1 * (1 - BM25_B + BM25_B * (chunkLength / avgdl));

    score += idf * (numerator / denom);
  }

  return score;
}

// ---------------------------------------------------------------------------
// Main strategy
// ---------------------------------------------------------------------------

/**
 * Score chunks by BM25-style relevance against a query.
 *
 * Algorithm:
 * 1. Tokenise the query into search terms.
 * 2. Split text into fixed-size chunks using the length-based strategy.
 * 3. Compute BM25 score for each chunk against the query terms.
 * 4. Return chunks sorted by relevance score descending.
 *
 * @param text - The input text to sample.
 * @param config - Sampling configuration (must include query for relevance).
 * @returns Chunks sorted by relevance score descending.
 * @throws If config.query is missing or empty.
 */
export function sampleByRelevance(
  text: string,
  config: SamplingConfig,
): readonly Chunk[] {
  if (text.length === 0) {
    return [];
  }

  const query = config.query?.trim();
  if (!query) {
    throw new Error(
      "The 'query' parameter is required when using the 'relevance' strategy.",
    );
  }

  const queryTokens = tokenise(query);
  if (queryTokens.length === 0) {
    throw new Error(
      "The 'query' must contain at least one meaningful word (2+ characters) " +
        "when using the 'relevance' strategy.",
    );
  }

  // Use length-based strategy to produce base chunks.
  const baseChunks = sampleByLength(text, config);

  if (baseChunks.length === 0) {
    return [];
  }

  // Precompute token lists for each chunk for DF computation.
  const chunkTokenLists = baseChunks.map((c) => tokenise(c.text));
  const df = documentFrequencies(chunkTokenLists);

  const avgChunkLength =
    baseChunks.reduce((sum, c) => sum + c.text.length, 0) / baseChunks.length;
  const avgdl = Math.max(avgChunkLength, AVG_CHUNK_LENGTH);

  // Score each chunk.
  const scored = baseChunks.map((chunk) => {
    const score = bm25Score(
      chunk.text,
      chunk.text.length,
      queryTokens,
      df,
      baseChunks.length,
      avgdl,
    );

    return {
      ...chunk,
      score,
    };
  });

  // Sort by score descending (most relevant first).
  // Stable sort preserves document order for equal scores.
  return scored.slice().sort((a, b) => {
    const diff = b.score - a.score;
    return diff !== 0 ? Math.sign(diff) : a.index - b.index;
  });
}