Blame for src/utils/indexingEngine.ts

import type {
  CorpusDocument,
  TermStats,
  DocIndexMeta,
  LengthHistogramBin,
  IndexTelemetry,
  BuiltInvertedIndex,
} from '../types/corpus';
import { computeEmbedding, computeSpladeTokens } from './retrievalScoring';

// Standard Lucene / English Stopwords
export const STANDARD_STOPWORDS = new Set([
  'a', 'about', 'above', 'after', 'again', 'against', 'all', 'am', 'an', 'and', 'any',
  'are', 'aren\'t', 'as', 'at', 'be', 'because', 'been', 'before', 'being', 'below',
  'between', 'both', 'but', 'by', 'can', 'cannot', 'could', 'did', 'do', 'does',
  'doing', 'don\'t', 'down', 'during', 'each', 'few', 'for', 'from', 'further',
  'had', 'has', 'have', 'having', 'he', 'her', 'here', 'hers', 'herself', 'him',
  'himself', 'his', 'how', 'i', 'if', 'in', 'into', 'is', 'it', 'its', 'itself',
  'me', 'more', 'most', 'my', 'myself', 'no', 'nor', 'not', 'of', 'off', 'on',
  'once', 'only', 'or', 'other', 'ought', 'our', 'ours', 'ourselves', 'out', 'over',
  'own', 'same', 'she', 'should', 'so', 'some', 'such', 'than', 'that', 'the',
  'their', 'theirs', 'them', 'themselves', 'then', 'there', 'these', 'they', 'this',
  'those', 'through', 'to', 'too', 'under', 'until', 'up', 'very', 'was', 'we',
  'were', 'what', 'when', 'where', 'which', 'while', 'who', 'whom', 'why', 'with',
  'would', 'you', 'your', 'yours', 'yourself', 'yourselves'
]);

export interface TokenWithPos {
  term: string;
  position: number;
}

// Tokenize text while tracking word position offsets
export function tokenizeWithPositions(
  text: string,
  filterStopwords: boolean = false
): TokenWithPos[] {
  const result: TokenWithPos[] = [];
  const regex = /[a-zA-Z0-9]+(?:'[a-zA-Z0-9]+)?/g;
  let match: RegExpExecArray | null;
  let pos = 0;

  while ((match = regex.exec(text)) !== null) {
    const term = match[0].toLowerCase();
    if (!filterStopwords || !STANDARD_STOPWORDS.has(term)) {
      result.push({ term, position: pos });
    }
    pos++;
  }

  return result;
}

// Compute length histogram bins
export function computeLengthHistogram(docLengths: { id: string; length: number }[]): LengthHistogramBin[] {
  if (docLengths.length === 0) return [];

  const lengths = docLengths.map((d) => d.length);
  const min = Math.min(...lengths);
  const max = Math.max(...lengths);

  if (min === max) {
    return [
      {
        binLabel: `${min} tokens`,
        rangeMin: min,
        rangeMax: max,
        count: docLengths.length,
        percentage: 100,
        docIds: docLengths.map((d) => d.id),
      },
    ];
  }

  const numBins = Math.min(6, Math.max(3, Math.ceil(Math.sqrt(docLengths.length))));
  const binWidth = Math.max(1, Math.ceil((max - min + 1) / numBins));

  const bins: LengthHistogramBin[] = [];
  for (let i = 0; i < numBins; i++) {
    const rangeMin = min + i * binWidth;
    const rangeMax = i === numBins - 1 ? max : rangeMin + binWidth - 1;
    bins.push({
      binLabel: `${rangeMin} - ${rangeMax}`,
      rangeMin,
      rangeMax,
      count: 0,
      percentage: 0,
      docIds: [],
    });
  }

  docLengths.forEach(({ id, length }) => {
    let placed = false;
    for (const bin of bins) {
      if (length >= bin.rangeMin && length <= bin.rangeMax) {
        bin.count++;
        bin.docIds.push(id);
        placed = true;
        break;
      }
    }
    if (!placed && bins.length > 0) {
      bins[bins.length - 1].count++;
      bins[bins.length - 1].docIds.push(id);
    }
  });

  const total = docLengths.length;
  bins.forEach((b) => {
    b.percentage = Number(((b.count / total) * 100).toFixed(1));
  });

  return bins.filter((b) => b.count > 0);
}

export interface IndexBuildOptions {
  filterStopwords?: boolean;
}

/**
 * Builds a complete client-side inverted index for a collection of documents.
 * Emulates Apache Lucene's core in-memory index structures and telemetry.
 */
export function buildInvertedIndex(
  documents: CorpusDocument[],
  options: IndexBuildOptions = {}
): BuiltInvertedIndex {
  const startTime = performance.now();
  const filterStopwords = options.filterStopwords ?? false;

  const N = documents.length;
  const dictionaryMap = new Map<string, TermStats>();
  const docMap = new Map<string, DocIndexMeta>();
  const docIndexList: DocIndexMeta[] = [];

  let totalTokens = 0;
  let rawSizeBytes = 0;

  // 1. Process documents and construct postings
  documents.forEach((doc, internalId) => {
    const fullText = (doc.title ? `${doc.title} ` : '') + doc.text;
    rawSizeBytes += new TextEncoder().encode(fullText).length;

    const tokensWithPos = tokenizeWithPositions(fullText, filterStopwords);
    const docLength = tokensWithPos.length;
    totalTokens += docLength;

    // Per-document term frequency and positions
    const docTfMap = new Map<string, { count: number; positions: number[] }>();
    tokensWithPos.forEach(({ term, position }) => {
      let entry = docTfMap.get(term);
      if (!entry) {
        entry = { count: 0, positions: [] };
        docTfMap.set(term, entry);
      }
      entry.count++;
      entry.positions.push(position);
    });

    // Populate global dictionary & inverted posting lists
    docTfMap.forEach((tfData, term) => {
      let termStat = dictionaryMap.get(term);
      if (!termStat) {
        termStat = {
          term,
          docFreq: 0,
          collectionFreq: 0,
          idf: 0,
          postings: [],
        };
        dictionaryMap.set(term, termStat);
      }

      termStat.docFreq++;
      termStat.collectionFreq += tfData.count;
      termStat.postings.push({
        docId: doc.id,
        docIndex: internalId,
        termFreq: tfData.count,
        positions: tfData.positions,
        docLength,
      });
    });

    // Term vector for document
    const termVector: { term: string; tf: number }[] = [];
    docTfMap.forEach((data, term) => {
      termVector.push({ term, tf: data.count });
    });
    termVector.sort((a, b) => b.tf - a.tf);

    // Compute dense embedding & SPLADE activations
    const denseVector = computeEmbedding(fullText, 768);
    const spladeWeights = computeSpladeTokens(fullText);
    const spladeExpansions = Object.entries(spladeWeights).map(([token, weight]) => ({
      token,
      weight,
      is_expansion: !fullText.toLowerCase().includes(token.toLowerCase()),
    }));

    const docMeta: DocIndexMeta = {
      id: doc.id,
      internalId,
      title: doc.title,
      text: doc.text,
      length: docLength,
      uniqueTerms: docTfMap.size,
      termVector,
      denseVector,
      denseDim: 768,
      denseSource: 'in-browser-simulation',
      denseLatencyMs: 0.1,
      spladeWeights,
      spladeExpansions,
      spladeSource: 'in-browser-simulation',
      spladeLatencyMs: 0.1,
      totalActivatedSplade: Object.keys(spladeWeights).length,
      mlmExpansionsCount: spladeExpansions.filter((e) => e.is_expansion).length,
    };

    docMap.set(doc.id, docMeta);
    docIndexList.push(docMeta);
  });

  // 2. Compute Lucene Robertson-Spärck Jones IDF for every term
  // IDF formula: ln(1 + (N - DF + 0.5) / (DF + 0.5))
  dictionaryMap.forEach((termStat) => {
    const df = termStat.docFreq;
    const rawIdf = Math.log(1.0 + (N - df + 0.5) / (df + 0.5));
    termStat.idf = Number(Math.max(0.1, rawIdf).toFixed(4));
    // Sort postings by Lucene internal Doc ID ascending (standard posting list invariant)
    termStat.postings.sort((a, b) => a.docIndex - b.docIndex);
  });

  const termsList: TermStats[] = Array.from(dictionaryMap.values());
  // Default sort by Document Frequency descending, then Collection Frequency
  termsList.sort((a, b) => {
    if (b.docFreq !== a.docFreq) return b.docFreq - a.docFreq;
    return b.collectionFreq - a.collectionFreq;
  });

  const endTime = performance.now();
  const indexingTimeMs = Number((endTime - startTime).toFixed(2));

  // 3. Compute telemetry metrics
  const avgdl = N > 0 ? Number((totalTokens / N).toFixed(2)) : 0;
  const vocabSize = dictionaryMap.size;

  const docLengths = docIndexList.map((d) => ({ id: d.id, length: d.length }));
  docLengths.sort((a, b) => a.length - b.length);

  const shortestDoc = docLengths[0] || { id: '-', length: 0 };
  const longestDoc = docLengths[docLengths.length - 1] || { id: '-', length: 0 };
  const histogram = computeLengthHistogram(docLengths);

  // Index size estimate:
  // Term dictionary: ~16 bytes overhead per term + string length
  // Postings: 4 bytes docDelta + 2 bytes tf + 2 bytes * positions
  let indexSizeBytes = 0;
  termsList.forEach((t) => {
    indexSizeBytes += 16 + t.term.length; // dictionary entry
    t.postings.forEach((p) => {
      indexSizeBytes += 4 + 2 + (p.positions ? p.positions.length * 2 : 0);
    });
  });

  const compressionRatio =
    indexSizeBytes > 0 ? Number((rawSizeBytes / indexSizeBytes).toFixed(2)) : 1.0;

  const telemetry: IndexTelemetry = {
    docCount: N,
    totalTokens,
    vocabSize,
    avgdl,
    indexingTimeMs,
    rawSizeBytes,
    indexSizeBytes,
    compressionRatio,
    shortestDoc,
    longestDoc,
    histogram,
  };

  return {
    telemetry,
    terms: termsList,
    dictionary: dictionaryMap,
    documents: docIndexList,
    docMap,
  };
}

// Term search filter (prefix or contains)
export function searchIndexTerms(
  index: BuiltInvertedIndex,
  query: string,
  mode: 'prefix' | 'contains' = 'contains'
): TermStats[] {
  if (!query.trim()) return index.terms;
  const q = query.trim().toLowerCase();

  return index.terms.filter((t) => {
    if (mode === 'prefix') {
      return t.term.startsWith(q);
    }
    return t.term.includes(q);
  });
}