data-sampler

/**
 * @file priority.ts
 * @description Priority-based sampling strategy.
 * Scores chunks by keyword frequency (auto-extracted from the text), position
 * (earlier content slightly boosted), and chunk density. Returns chunks sorted
 * by score descending.
 */

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

/**
 * A token with its frequency across the document.
 */
interface TokenFrequency {
  readonly token: string;
  readonly count: number;
}

/**
 * Minimum character length for a word to be considered a keyword.
 * Filters out short connector words, articles, and very short tokens.
 */
const MIN_KEYWORD_LENGTH = 4;

/**
 * Maximum number of top keywords to use for scoring.
 */
const MAX_KEYWORDS = 20;

/**
 * Score boost factor for positional proximity to the start of the text.
 * 0 = no boost, higher = stronger early-position bias.
 */
const POSITION_BOOST_FACTOR = 0.15;

// ---------------------------------------------------------------------------
// Keyword extraction
// ---------------------------------------------------------------------------

/**
 * Extract the most frequent meaningful keywords from text.
 * Filters out short words, normalises case, and returns the top N tokens
 * by descending frequency.
 */
function extractKeywords(text: string): readonly TokenFrequency[] {
  const words = text.toLowerCase().match(/\p{L}{4,}/gu) ?? [];

  const freq = new Map<string, number>();
  for (const word of words) {
    freq.set(word, (freq.get(word) ?? 0) + 1);
  }

  const sorted = [...freq.entries()]
    .map(([token, count]) => ({ token, count }))
    .sort((a, b) => b.count - a.count);

  return sorted.slice(0, MAX_KEYWORDS);
}

// ---------------------------------------------------------------------------
// Scoring
// ---------------------------------------------------------------------------

/**
 * Compute a keyword density score for a chunk.
 * Returns the fraction of keyword occurrences that fall within this chunk
 * relative to the total keyword occurrences in the whole text.
 *
 * @param chunkText - The text of the chunk.
 * @param keywords - The top keywords extracted from the full text.
 * @returns A density score between 0 and 1.
 */
function keywordDensityScore(
  chunkText: string,
  keywords: readonly TokenFrequency[],
): number {
  if (keywords.length === 0) return 0;

  const lower = chunkText.toLowerCase();
  let score = 0;

  for (const kw of keywords) {
    const re = new RegExp(
      kw.token.replace(/[.*+?^${}()|[\]\\]/gu, "\\$&"),
      "gu",
    );
    const matches = lower.match(re);
    if (matches) {
      // Each keyword contributes proportionally to its global frequency.
      score += matches.length / kw.count;
    }
  }

  return Math.min(score / keywords.length, 1);
}

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

/**
 * Score chunks by keyword frequency and position.
 *
 * Algorithm:
 * 1. Extract top keywords from the full text (by frequency).
 * 2. Split text into fixed-size chunks using the length-based strategy.
 * 3. Score each chunk by keyword density (primary) and position (secondary boost).
 * 4. Return chunks sorted by score descending (highest priority first).
 *
 * @param text - The input text to sample.
 * @param config - Sampling configuration (chunkSize, overlap, maxChunks).
 * @returns Chunks sorted by priority score descending.
 */
export function sampleByPriority(
  text: string,
  config: SamplingConfig,
): readonly Chunk[] {
  if (text.length === 0) {
    return [];
  }

  const keywords = extractKeywords(text);

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

  // Score each chunk.
  const scored = baseChunks.map((chunk) => {
    const density = keywordDensityScore(chunk.text, keywords);

    // Position boost: earlier chunks get a small bonus.
    // Normalised to 0-1 based on chunk index relative to total.
    const positionRatio =
      baseChunks.length > 1 ? 1 - chunk.index / (baseChunks.length - 1) : 1;
    const positionBoost = positionRatio * POSITION_BOOST_FACTOR;

    // Combined score: density dominates, position is a tiebreaker.
    const score = Math.min(density + positionBoost, 1);

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

  // Sort by score descending (highest priority 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;
  });
}