/**
 * Small dense-matrix helpers used by the LSA (truncated SVD) step.
 *
 * The corpus is ~40 documents, so the whole decomposition fits comfortably in
 * memory as plain `number[][]`. An exact symmetric eigendecomposition is both
 * simpler and more accurate here than a randomised approximation would be.
 */

export type Matrix = number[][];

export function zeros(rows: number, cols: number): Matrix {
  return Array.from({ length: rows }, () => new Array<number>(cols).fill(0));
}

export function identity(size: number): Matrix {
  const result = zeros(size, size);
  for (let i = 0; i < size; i += 1) result[i][i] = 1;
  return result;
}

/** Matrix product `a · b`. */
export function matMul(a: Matrix, b: Matrix): Matrix {
  const rows = a.length;
  const inner = b.length;
  const cols = b[0]?.length ?? 0;
  const result = zeros(rows, cols);

  for (let i = 0; i < rows; i += 1) {
    const rowA = a[i];
    const rowOut = result[i];
    for (let k = 0; k < inner; k += 1) {
      const coeff = rowA[k];
      if (coeff === 0) continue;
      const rowB = b[k];
      for (let j = 0; j < cols; j += 1) {
        rowOut[j] += coeff * rowB[j];
      }
    }
  }

  return result;
}

export function transpose(matrix: Matrix): Matrix {
  const rows = matrix.length;
  const cols = matrix[0]?.length ?? 0;
  const result = zeros(cols, rows);
  for (let i = 0; i < rows; i += 1) {
    for (let j = 0; j < cols; j += 1) result[j][i] = matrix[i][j];
  }
  return result;
}

export function dot(a: readonly number[], b: readonly number[]): number {
  let total = 0;
  for (let i = 0; i < a.length; i += 1) total += a[i] * b[i];
  return total;
}

export function norm(vector: readonly number[]): number {
  return Math.sqrt(dot(vector, vector));
}

/** Returns a unit-length copy; a zero vector is returned unchanged. */
export function normalize(vector: readonly number[]): number[] {
  const length = norm(vector);
  if (length === 0) return [...vector];
  return vector.map((value) => value / length);
}

export interface SymmetricEigen {
  /** Eigenvalues, descending. */
  values: number[];
  /** Column `i` of `vectors` is the eigenvector for `values[i]`. */
  vectors: Matrix;
}

/**
 * Cyclic Jacobi eigendecomposition of a real symmetric matrix.
 *
 * `A = V · diag(values) · Vᵀ`. Converges quadratically and is exact to
 * floating-point tolerance for the small (n <= a few hundred) matrices used
 * here, which is why no randomised range finder is needed.
 */
export function jacobiEigenSymmetric(
  input: Matrix,
  { maxSweeps = 60, tolerance = 1e-12 }: { maxSweeps?: number; tolerance?: number } = {},
): SymmetricEigen {
  const n = input.length;
  const a: Matrix = input.map((row) => [...row]);
  const vectors = identity(n);

  for (let sweep = 0; sweep < maxSweeps; sweep += 1) {
    let offDiagonal = 0;
    for (let p = 0; p < n; p += 1) {
      for (let q = p + 1; q < n; q += 1) offDiagonal += a[p][q] * a[p][q];
    }
    if (Math.sqrt(offDiagonal) < tolerance) break;

    for (let p = 0; p < n; p += 1) {
      for (let q = p + 1; q < n; q += 1) {
        const apq = a[p][q];
        if (Math.abs(apq) < 1e-15) continue;

        // Rotation angle chosen to annihilate a[p][q] in one sweep.
        const theta = (a[q][q] - a[p][p]) / (2 * apq);
        const sign = theta >= 0 ? 1 : -1;
        const t = sign / (Math.abs(theta) + Math.sqrt(theta * theta + 1));
        const c = 1 / Math.sqrt(t * t + 1);
        const s = t * c;

        for (let k = 0; k < n; k += 1) {
          const akp = a[k][p];
          const akq = a[k][q];
          a[k][p] = c * akp - s * akq;
          a[k][q] = s * akp + c * akq;
        }
        for (let k = 0; k < n; k += 1) {
          const apk = a[p][k];
          const aqk = a[q][k];
          a[p][k] = c * apk - s * aqk;
          a[q][k] = s * apk + c * aqk;
        }
        for (let k = 0; k < n; k += 1) {
          const vkp = vectors[k][p];
          const vkq = vectors[k][q];
          vectors[k][p] = c * vkp - s * vkq;
          vectors[k][q] = s * vkp + c * vkq;
        }
      }
    }
  }

  const order = Array.from({ length: n }, (_, index) => index).sort(
    (left, right) => a[right][right] - a[left][left],
  );

  return {
    values: order.map((index) => a[index][index]),
    vectors: Array.from({ length: n }, (_, row) => order.map((index) => vectors[row][index])),
  };
}