freezmo/lib/notes-source.ts

/**
 * "Where in my notes did this come from?" — the study feed shows the passage
 * that carries an activity's answer, with the matching words highlighted.
 *
 * Two paths, both local (nothing leaves the browser):
 *  - activities generated after the source-quote change carry a verbatim
 *    `source` quote; it is located in the stored notes and the sentence around
 *    it is shown, with the quote itself highlighted;
 *  - older activities have no quote, so the passage is found by scoring the
 *    notes' sentences against the activity's own wording. Rarer words weigh
 *    more (an IDF sum), so "Rayleigh" beats "the", and a sentence has to match
 *    at least two distinct terms to count.
 *
 * Returns null when nothing convincing is found — a random sentence presented
 * as "where this came from" would be worse than showing nothing.
 */

import type { Activity } from "./schemas";

export type NotesMatch = {
  /** The passage to show. */
  text: string;
  /** Highlight spans inside `text` (start, end). */
  spans: Array<[number, number]>;
  /** True when the AI supplied the quote (vs. a local sentence match). */
  quote: boolean;
};

/** Longest excerpt shown under an activity. */
const MAX_EXCERPT = 280;

const STOPWORDS = new Set([
  "the", "and", "that", "with", "from", "this", "which", "are", "was", "were", "have", "has", "had",
  "for", "not", "but", "you", "your", "they", "their", "into", "than", "then", "what", "when",
  "where", "why", "how", "does", "did", "can", "could", "should", "would", "its", "one", "two",
  "about", "over", "under", "more", "most", "some", "any", "all", "each", "also", "such", "only",
  "very", "just", "like", "make", "made", "use", "used", "using", "how", "who", "whose", "will",
]);

const isWordCharacter = (character: string | undefined) =>
  Boolean(character && /[\p{L}\p{N}]/u.test(character));

/** Lower-case the words of a text, for term comparison. */
export function termsOf(text: string): string[] {
  return [
    ...new Set(
      text
        .toLowerCase()
        .split(/[^\p{L}\p{N}]+/u)
        .filter((word) => word.length >= 3 && !STOPWORDS.has(word)),
    ),
  ];
}

/**
 * The activity's own wording, used when no quote is stored: the question side
 * plus the answer side, so both "what was asked" and "what the answer mentions"
 * can hit a sentence in the notes.
 */
export function activityQuery(activity: Activity): string {
  switch (activity.type) {
    case "flashcard":
      return `${activity.front} ${activity.back}`;
    case "multiple_choice":
      return `${activity.question} ${activity.options[activity.correctIndex] ?? ""}`;
    case "true_false":
      return activity.statement;
    case "fill_blank":
      return `${activity.sentence.replace(/_{2,}/g, " ")} ${activity.answer}`;
    case "matching":
      return activity.pairs.map((pair) => `${pair.term} ${pair.definition}`).join(" ");
    case "ordering":
      return activity.items.join(" ");
    case "short_answer":
      return `${activity.question} ${activity.referenceAnswer} ${activity.keyPoints.join(" ")}`;
    default:
      return "";
  }
}

/** `text` plus a per-character index map into the original. */
function normalise(text: string): { text: string; map: number[] } {
  let out = "";
  const map: number[] = [];
  let pendingSpace = false;
  for (let index = 0; index < text.length; index += 1) {
    const character = text[index];
    if (/\s/.test(character)) {
      pendingSpace = out.length > 0;
      continue;
    }
    if (pendingSpace) {
      out += " ";
      map.push(index);
      pendingSpace = false;
    }
    out += character.toLowerCase();
    map.push(index);
  }
  return { text: out, map };
}

/** The quote (normalised) inside the notes, or the longest verbatim word run. */
function locateQuote(notes: string, quote: string): [number, number] | null {
  const haystack = normalise(notes);
  const needle = normalise(quote).text.replace(/…/g, "").trim();
  if (needle.length < 8) return null;

  const span = (from: number, to: number): [number, number] => [
    haystack.map[from],
    haystack.map[to - 1] + 1,
  ];

  const exact = haystack.text.indexOf(needle);
  if (exact >= 0) return span(exact, exact + needle.length);

  // The model paraphrased or trimmed the quote: fall back to the longest run
  // of consecutive words that does appear verbatim.
  const words = needle.split(" ").filter(Boolean);
  for (let size = Math.min(words.length, 16); size >= 3; size -= 1) {
    for (let start = 0; start + size <= words.length; start += 1) {
      const chunk = words.slice(start, start + size).join(" ");
      const at = haystack.text.indexOf(chunk);
      if (at >= 0) return span(at, at + chunk.length);
    }
  }
  return null;
}

/** Expand a match to the sentence (or line) around it. */
function sentenceAround(text: string, start: number, end: number): [number, number] {
  let from = start;
  while (from > 0) {
    const previous = text[from - 1];
    if (previous === "\n" || previous === "." || previous === "!" || previous === "?") break;
    from -= 1;
  }
  let to = end;
  while (to < text.length) {
    const character = text[to];
    if (character === "\n" || character === "." || character === "!" || character === "?") {
      to += 1;
      break;
    }
    to += 1;
  }
  while (from < to && /\s/.test(text[from])) from += 1;
  while (to > from && /\s/.test(text[to - 1])) to -= 1;
  return [from, to];
}

/** Trim a passage to `max` characters around the first highlight. */
function clip(
  text: string,
  start: number,
  end: number,
  spans: Array<[number, number]>,
  max: number,
): { text: string; spans: Array<[number, number]> } | null {
  const length = end - start;
  if (length <= max) {
    return {
      text: text.slice(start, end),
      spans: spans
        .map(([from, to]) => [Math.max(start, from), Math.min(end, to)] as [number, number])
        .filter(([from, to]) => to > from)
        .map(([from, to]) => [from - start, to - start] as [number, number]),
    };
  }
  const anchor = spans[0] ?? [start, start];
  const from = Math.max(start, Math.min(anchor[0] - Math.floor(max / 3), end - max));
  const to = Math.min(end, from + max);
  const prefix = from > start ? "… " : "";
  const suffix = to < end ? " …" : "";
  const trimmed = `${prefix}${text.slice(from, to).trim()}${suffix}`;
  const shift = from + prefix.length;
  return {
    text: trimmed,
    spans: spans
      .map(([spanFrom, spanTo]) => [Math.max(from, spanFrom) - shift, Math.min(to, spanTo) - shift] as [number, number])
      .filter(([spanFrom, spanTo]) => spanTo > spanFrom && spanFrom >= 0 && spanTo <= trimmed.length),
  };
}

/** Sentence split that keeps offsets; no lookbehind (older Safari). */
function sentences(text: string): Array<{ start: number; end: number }> {
  const out: Array<{ start: number; end: number }> = [];
  let start = 0;
  const push = (end: number) => {
    let from = start;
    while (from < end && /\s/.test(text[from])) from += 1;
    let to = end;
    while (to > from && /\s/.test(text[to - 1])) to -= 1;
    if (to - from >= 8) out.push({ start: from, end: to });
    start = end;
  };
  for (let index = 0; index < text.length; index += 1) {
    const character = text[index];
    if (character === "\n") {
      push(index);
      start = index + 1;
      continue;
    }
    if ((character === "." || character === "!" || character === "?") && /\s/.test(text[index + 1] ?? "\n")) {
      push(index + 1);
      start = index + 1;
    }
  }
  push(text.length);
  return out;
}

/** Every occurrence of the matched terms in a sentence (word boundaries). */
function termSpans(text: string, matched: Set<string>): Array<[number, number]> {
  const { text: haystack, map } = normalise(text);
  const spans: Array<[number, number]> = [];
  for (const term of matched) {
    let at = haystack.indexOf(term);
    while (at >= 0) {
      if (!isWordCharacter(haystack[at - 1]) && !isWordCharacter(haystack[at + term.length])) {
        spans.push([map[at], map[at + term.length - 1] + 1]);
      }
      at = haystack.indexOf(term, at + 1);
    }
  }
  return spans.sort((a, b) => a[0] - b[0]);
}

/** The best-matching sentence for a query, IDF-weighted. */
export function bestSentence(notes: string, query: string): NotesMatch | null {
  const list = sentences(notes);
  if (list.length === 0) return null;
  const queryTerms = termsOf(query);
  if (queryTerms.length === 0) return null;

  const perSentence = list.map((sentence) => new Set(termsOf(notes.slice(sentence.start, sentence.end))));
  const documentFrequency = new Map<string, number>();
  for (const unique of perSentence) {
    for (const term of unique) documentFrequency.set(term, (documentFrequency.get(term) ?? 0) + 1);
  }
  const total = list.length;

  let best: { index: number; score: number; matched: Set<string> } | null = null;
  perSentence.forEach((unique, index) => {
    const matched = new Set<string>();
    let score = 0;
    for (const term of queryTerms) {
      if (!unique.has(term)) continue;
      matched.add(term);
      score += Math.log(1 + total / (documentFrequency.get(term) ?? 1));
    }
    const enough = matched.size >= 2 || (matched.size === 1 && total <= 3);
    if (!enough) return;
    if (!best || score > best.score) best = { index, score, matched };
  });
  if (!best) return null;

  const winner: { index: number; score: number; matched: Set<string> } = best;
  const sentence = list[winner.index];
  const spans = termSpans(notes.slice(sentence.start, sentence.end), winner.matched);
  const clipped = clip(notes, sentence.start, sentence.end, spans, MAX_EXCERPT);
  if (!clipped) return null;
  return { text: clipped.text, spans: clipped.spans, quote: false };
}

/**
 * The passage that carries this activity's answer, from the AI's own quote
 * when there is one and by local matching otherwise.
 */
export function notesSource(notes: string | undefined, activity: Activity): NotesMatch | null {
  const text = (notes ?? "").trim();
  if (text.length < 24) return null;

  if (activity.source) {
    const located = locateQuote(text, activity.source);
    if (located) {
      const [start, end] = sentenceAround(text, located[0], located[1]);
      const clipped = clip(text, start, end, [located], MAX_EXCERPT);
      if (clipped) return { text: clipped.text, spans: clipped.spans, quote: true };
    }
  }
  return bestSentence(text, activityQuery(activity));
}

/** Splits a passage into pieces for rendering, marking the highlighted ones. */
export function highlightSegments(
  text: string,
  spans: Array<[number, number]>,
): Array<{ text: string; hit: boolean }> {
  const merged = [...spans]
    .filter(([from, to]) => to > from && from >= 0 && to <= text.length)
    .sort((a, b) => a[0] - b[0])
    .reduce<Array<[number, number]>>((list, span) => {
      const last = list[list.length - 1];
      if (last && span[0] <= last[1]) last[1] = Math.max(last[1], span[1]);
      else list.push([span[0], span[1]]);
      return list;
    }, []);
  if (merged.length === 0) return [{ text, hit: false }];
  const out: Array<{ text: string; hit: boolean }> = [];
  let at = 0;
  for (const [from, to] of merged) {
    if (from > at) out.push({ text: text.slice(at, from), hit: false });
    out.push({ text: text.slice(from, to), hit: true });
    at = to;
  }
  if (at < text.length) out.push({ text: text.slice(at), hit: false });
  return out;
}