import { tokenize } from "@/lib/text";
import { type Matrix, jacobiEigenSymmetric, matMul, normalize, transpose, zeros } from "@/lib/matrix";

/**
 * Latent Semantic Analysis over the product corpus.
 *
 * This is the offline embedding provider: instead of calling a neural model, it
 * factorises the product term-document matrix into a shared latent space in
 * which both documents and *new, unseen* text can be represented. That is what
 * lets a query for "headphones for travelling" reach a product described as
 * "noise cancelling, long flights" even though the two share no single term.
 *
 *   X  (documents x terms)  -- tf-idf, L2 normalised rows
 *   G  = X · Xᵀ            -- Gram matrix
 *   G  = U · Σ² · Uᵀ       -- exact symmetric eigendecomposition
 *   document vectors D = U · Σ
 *   term vectors     T = Xᵀ · U · Σ⁻¹      (T has orthonormal columns)
 *
 * Folding a query in is a weighted sum of its term vectors, which means
 * `cos(D_i, q)` reduces to a projected tf-idf score — the classic LSA
 * pseudo-relevance formulation.
 */

export interface LsaModel {
  /** Vocabulary, aligned with every row of `termVectors`. */
  terms: string[];
  /** Inverse document frequency per term, aligned with `terms`. */
  idf: number[];
  /** `terms.length x dimensions`, columns orthonormal. */
  termVectors: Matrix;
  /** Number of latent dimensions retained. */
  dimensions: number;
  /** Singular values, descending. */
  singularValues: number[];
}

export interface LsaCorpus {
  model: LsaModel;
  /** One unit-length vector per input document, in input order. */
  documentVectors: number[][];
}

const DEFAULT_DIMENSIONS = 24;
/** Terms must appear in at least this many documents to enter the vocabulary. */
const MIN_DOCUMENT_FREQUENCY = 1;

function buildTermDocumentMatrix(
  documents: string[],
): { matrix: Matrix; terms: string[]; idf: number[] } {
  const documentCount = documents.length;
  const termFrequencies: Map<string, number>[] = [];
  const documentFrequency = new Map<string, number>();

  for (const document of documents) {
    const tokens = tokenize(document);
    const counts = new Map<string, number>();
    for (const token of tokens) {
      counts.set(token, (counts.get(token) ?? 0) + 1);
    }
    termFrequencies.push(counts);
    for (const term of counts.keys()) {
      documentFrequency.set(term, (documentFrequency.get(term) ?? 0) + 1);
    }
  }

  const terms = [...documentFrequency.entries()]
    .filter(([, frequency]) => frequency >= MIN_DOCUMENT_FREQUENCY)
    .map(([term]) => term)
    .sort();

  const termIndex = new Map(terms.map((term, index) => [term, index]));
  const idf = terms.map((term) => {
    const frequency = documentFrequency.get(term) ?? 1;
    // Smoothed idf, floored so that a term present in every document still
    // contributes a little to the query representation.
    return Math.max(0.05, Math.log(1 + documentCount / frequency));
  });

  const matrix = zeros(documentCount, terms.length);

  for (let row = 0; row < documentCount; row += 1) {
    const counts = termFrequencies[row];
    for (const [term, frequency] of counts) {
      const column = termIndex.get(term);
      if (column === undefined) continue;
      // Sublinear tf keeps a tag repeated in a long description from dominating.
      matrix[row][column] = (1 + Math.log(frequency)) * idf[column];
    }
    const rowVector = matrix[row];
    const length = Math.sqrt(rowVector.reduce((total, value) => total + value * value, 0));
    if (length > 0) {
      for (let column = 0; column < rowVector.length; column += 1) rowVector[column] /= length;
    }
  }

  return { matrix, terms, idf };
}

/**
 * Factorise the corpus into an LSA model plus one vector per document.
 *
 * `dimensions` is clamped to the corpus rank, so a 3-document corpus still
 * produces a valid model instead of returning null directions.
 */
export function buildLsaCorpus(
  documents: string[],
  requestedDimensions: number = DEFAULT_DIMENSIONS,
): LsaCorpus {
  const { matrix, terms, idf } = buildTermDocumentMatrix(documents);
  const documentCount = documents.length;
  const termCount = terms.length;

  const maxDimensions = Math.max(1, Math.min(documentCount, termCount) - 1);
  const dimensions = Math.max(1, Math.min(requestedDimensions, maxDimensions));

  // Gram matrix of the tf-idf matrix, then its symmetric eigendecomposition.
  const gram = matMul(matrix, transpose(matrix));
  const { values, vectors } = jacobiEigenSymmetric(gram);

  const singularValues: number[] = [];
  const docBasis: Matrix = [];
  for (let index = 0; index < dimensions; index += 1) {
    // Negative values are pure floating-point noise on a PSD matrix.
    const singularValue = Math.sqrt(Math.max(0, values[index]));
    singularValues.push(singularValue);
    docBasis.push(vectors.map((row) => row[index]));
  }

  // documentVectors = U · Σ, one unit-length row per document.
  const documentVectors = matrix.map((_, rowIndex) => {
    const projection = new Array<number>(dimensions).fill(0);
    for (let component = 0; component < dimensions; component += 1) {
      projection[component] = vectors[rowIndex][component] * singularValues[component];
    }
    return normalize(projection);
  });

  // termVectors = Xᵀ · U · Σ⁻¹ — columns are orthonormal, so term vectors can be
  // summed to represent text that never appeared in the corpus.
  const termVectors: Matrix = [];
  for (let column = 0; column < termCount; column += 1) {
    const projection = new Array<number>(dimensions).fill(0);
    for (let component = 0; component < dimensions; component += 1) {
      const singularValue = singularValues[component];
      if (singularValue <= 1e-9) continue;
      let sum = 0;
      for (let row = 0; row < documentCount; row += 1) {
        sum += matrix[row][column] * vectors[row][component];
      }
      projection[component] = sum / singularValue;
    }
    termVectors.push(normalize(projection));
  }

  return {
    model: { terms, idf, termVectors, dimensions, singularValues },
    documentVectors,
  };
}

/** Fold arbitrary text into the latent space as a weighted sum of term vectors. */
export function projectQuery(model: LsaModel, text: string): number[] {
  const termIndex = new Map(model.terms.map((term, index) => [term, index]));
  const tokens = tokenize(text);
  const weights = new Map<string, number>();

  for (const token of tokens) {
    weights.set(token, (weights.get(token) ?? 0) + 1);
  }

  const vector = new Array<number>(model.dimensions).fill(0);
  for (const [token, occurrences] of weights) {
    const column = termIndex.get(token);
    if (column === undefined) continue;
    const termVector = model.termVectors[column];
    // idf keeps common words from swamping the rare, discriminative ones.
    const weight = (1 + Math.log(occurrences)) * model.idf[column];
    for (let component = 0; component < model.dimensions; component += 1) {
      vector[component] += weight * termVector[component];
    }
  }

  return normalize(vector);
}