// Shift-invert Lanczos for the eigenvalues of the pencil (K, M) in one
// slice [lo, hi] of the spectrum. With M diagonal the pencil is turned into
// the standard symmetric problem A = M^{-1/2} K M^{-1/2}, and Lanczos is
// run on (A - sigma I)^{-1} = M^{1/2} (K - sigma M)^{-1} M^{1/2} with the
// banded factorization of banded.ts, full reorthogonalization, and a check
// of the count of converged eigenvalues in the slice against the inertia
// counts at its ends. The mode shapes returned are M^{-1/2} times the
// eigenvectors of A, that is, mass-normalized.

import { factorize, solve, type Factorization } from "./banded.ts";

export interface SliceResult {
  /** Eigenvalues omega^2 found in (lo, hi], ascending. */
  lambda: number[];
  /** Mass-normalized mode shapes, one Float64Array of length n per eigenvalue. */
  shapes: Float64Array[];
  /** Number of eigenvalues the inertia counts say lie in the slice. */
  expected: number;
  /** Lanczos steps used. */
  steps: number;
}

export interface SliceOptions {
  /** Relative tolerance on the shift-inverted Ritz values; 1e-9 by default. */
  tol?: number;
  /** Largest number of Lanczos steps per start; 3 count + 40 by default. */
  maxSteps?: number;
  /** Restarts with a fresh random vector when modes are missing (degenerate pairs); 3 by default. */
  restarts?: number;
  seed?: number;
}

/**
 * Eigenvalues of (K, M) in (lo, hi] and their shapes. `countBelowLo` and
 * `countBelowHi` are the inertia counts at the two ends, so that the
 * expected number of eigenvalues is known.
 */
export function sliceModes(
  band: Float64Array,
  mass: Float64Array,
  n: number,
  bw: number,
  lo: number,
  hi: number,
  countBelowLo: number,
  countBelowHi: number,
  opts: SliceOptions = {},
): SliceResult {
  const expected = countBelowHi - countBelowLo;
  if (expected <= 0) return { lambda: [], shapes: [], expected, steps: 0 };
  const tol = opts.tol ?? 1e-9;
  const restarts = opts.restarts ?? 3;
  const sigma = 0.5 * (lo + hi);
  const fact = factorize(band, mass, n, bw, sigma);
  const sqm = new Float64Array(n);
  for (let i = 0; i < n; i++) sqm[i] = Math.sqrt(mass[i]);
  let seed = (opts.seed ?? 4321) >>> 0;
  const rnd = () => {
    seed = (seed * 1664525 + 1013904223) >>> 0;
    return seed / 4294967296 - 0.5;
  };

  const locked: { lambda: number; x: Float64Array }[] = [];
  let steps = 0;
  const maxSteps = Math.min(n, opts.maxSteps ?? 3 * expected + 40);
  for (let attempt = 0; attempt <= restarts && locked.length < expected; attempt++) {
    const found = lanczosRun(fact, sqm, n, sigma, lo, hi, expected - locked.length, locked.map((l) => l.x), maxSteps, tol, rnd);
    steps += found.steps;
    for (const f of found.pairs) locked.push(f);
  }
  locked.sort((a, b) => a.lambda - b.lambda);
  const shapes = locked.map((l) => {
    const phi = new Float64Array(n);
    for (let i = 0; i < n; i++) phi[i] = l.x[i] / sqm[i];
    return phi;
  });
  return { lambda: locked.map((l) => l.lambda), shapes, expected, steps };
}

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

/** One Lanczos run from a random start, orthogonal to the locked vectors; returns the converged pairs in (lo, hi]. */
function lanczosRun(
  fact: Factorization,
  sqm: Float64Array,
  n: number,
  sigma: number,
  lo: number,
  hi: number,
  wanted: number,
  locked: Float64Array[],
  maxSteps: number,
  tol: number,
  rnd: () => number,
): { pairs: { lambda: number; x: Float64Array }[]; steps: number } {
  const Q: Float64Array[] = [];
  const alpha: number[] = [], beta: number[] = [];
  const tmp = new Float64Array(n);
  // Classical Gram-Schmidt against the basis, repeated once when the first
  // pass removed a large part of the vector ("twice is enough"); the second
  // pass is rarely needed and the basis is the main memory traffic.
  const orthogonalize = (r: Float64Array, against: Float64Array[]): void => {
    if (against.length === 0) return;
    const before = dot(r, r);
    for (let pass = 0; pass < 2; pass++) {
      for (const q of against) {
        const c = dot(r, q);
        if (c !== 0) for (let i = 0; i < n; i++) r[i] -= c * q[i];
      }
      if (dot(r, r) > 0.5 * before) break;
    }
  };
  let q = new Float64Array(n);
  for (let i = 0; i < n; i++) q[i] = rnd();
  orthogonalize(q, locked);
  let nrm = Math.sqrt(dot(q, q));
  for (let i = 0; i < n; i++) q[i] /= nrm;
  Q.push(q);

  let pairs: { lambda: number; x: Float64Array }[] = [];
  let steps = 0;
  const checkEvery = Math.max(6, Math.round(wanted / 10));
  const check = (): boolean => {
    const m = Q.length;
    const { values, vectors } = tridiagonalEigen(alpha, beta.slice(0, m - 1));
    const betaM = beta[m - 1] ?? 0;
    const conv: { lambda: number; s: Float64Array }[] = [];
    for (let k = 0; k < m; k++) {
      const theta = values[k];
      if (theta === 0) continue;
      const lambda = sigma + 1 / theta;
      if (!(lambda > lo && lambda <= hi)) continue;
      const err = Math.abs(betaM * vectors[k][m - 1]);
      if (err <= tol * Math.abs(theta)) conv.push({ lambda, s: vectors[k] });
    }
    if (conv.length < wanted && m < maxSteps) return false;
    pairs = conv.map(({ lambda, s }) => {
      const x = new Float64Array(n);
      for (let j = 0; j < m; j++) {
        const sj = s[j], qj = Q[j];
        if (sj === 0) continue;
        for (let i = 0; i < n; i++) x[i] += sj * qj[i];
      }
      return { lambda, x };
    });
    return true;
  };

  for (let j = 0; j < maxSteps; j++) {
    steps++;
    // r = Op q_j with Op = M^{1/2} (K - sigma M)^{-1} M^{1/2}
    for (let i = 0; i < n; i++) tmp[i] = sqm[i] * Q[j][i];
    solve(fact, tmp, tmp);
    const r = new Float64Array(n);
    for (let i = 0; i < n; i++) r[i] = sqm[i] * tmp[i];
    if (j > 0) for (let i = 0; i < n; i++) r[i] -= beta[j - 1] * Q[j - 1][i];
    const a = dot(r, Q[j]);
    alpha.push(a);
    for (let i = 0; i < n; i++) r[i] -= a * Q[j][i];
    orthogonalize(r, Q);
    orthogonalize(r, locked);
    const b = Math.sqrt(dot(r, r));
    beta.push(b);
    const last = j === maxSteps - 1 || Q.length >= n;
    if (b < 1e-14 * Math.max(1, Math.abs(a))) {
      // Invariant subspace: everything reachable has converged.
      check();
      break;
    }
    // Fewer than `wanted` steps cannot hold `wanted` converged Ritz values,
    // and the tridiagonal eigenproblem costs m^3, so check only from there
    // on and at a spacing that grows with the slice.
    const m = j + 1;
    if (last || (m >= wanted && (m - wanted) % checkEvery === 0)) {
      if (check()) break;
    }
    if (Q.length >= n) break;
    for (let i = 0; i < n; i++) r[i] /= b;
    Q.push(r);
  }
  if (pairs.length === 0 && steps > 0) check();
  return { pairs, steps };
}

/**
 * Eigenvalues and eigenvectors of the symmetric tridiagonal matrix with
 * diagonal `d` and off-diagonal `e` (implicit QL, after EISPACK tql2).
 * vectors[k] is the eigenvector of values[k].
 */
export function tridiagonalEigen(dIn: number[], eIn: number[]): { values: number[]; vectors: Float64Array[] } {
  const n = dIn.length;
  const d = dIn.slice();
  const e = eIn.slice();
  e.push(0);
  const z: Float64Array[] = [];
  for (let i = 0; i < n; i++) {
    const row = new Float64Array(n);
    row[i] = 1;
    z.push(row);
  }
  let f = 0, tst1 = 0;
  const eps = 2 ** -52;
  for (let l = 0; l < n; l++) {
    tst1 = Math.max(tst1, Math.abs(d[l]) + Math.abs(e[l]));
    let m = l;
    while (m < n) {
      if (Math.abs(e[m]) <= eps * tst1) break;
      m++;
    }
    if (m > l) {
      let iter = 0;
      do {
        iter++;
        if (iter > 200) throw new Error("tridiagonal QL did not converge");
        let g = d[l];
        let p = (d[l + 1] - g) / (2 * e[l]);
        let r = Math.hypot(p, 1);
        if (p < 0) r = -r;
        d[l] = e[l] / (p + r);
        d[l + 1] = e[l] * (p + r);
        const dl1 = d[l + 1];
        let h = g - d[l];
        for (let i = l + 2; i < n; i++) d[i] -= h;
        f += h;
        p = d[m];
        let c = 1, c2 = c, c3 = c;
        const el1 = e[l + 1];
        let s = 0, s2 = 0;
        for (let i = m - 1; i >= l; i--) {
          c3 = c2; c2 = c; s2 = s;
          g = c * e[i];
          h = c * p;
          r = Math.hypot(p, e[i]);
          e[i + 1] = s * r;
          s = e[i] / r;
          c = p / r;
          p = c * d[i] - s * g;
          d[i + 1] = h + s * (c * g + s * d[i]);
          for (let k = 0; k < n; k++) {
            h = z[k][i + 1];
            z[k][i + 1] = s * z[k][i] + c * h;
            z[k][i] = c * z[k][i] - s * h;
          }
        }
        p = (-s * s2 * c3 * el1 * e[l]) / dl1;
        e[l] = s * p;
        d[l] = c * p;
      } while (Math.abs(e[l]) > eps * tst1);
    }
    d[l] += f;
    e[l] = 0;
  }
  const order = Array.from({ length: n }, (_, i) => i).sort((a, b) => d[a] - d[b]);
  const values = order.map((i) => d[i]);
  const vectors = order.map((i) => {
    const v = new Float64Array(n);
    for (let k = 0; k < n; k++) v[k] = z[k][i];
    return v;
  });
  return { values, vectors };
}