// The modes of a plate up to a cutoff frequency, by spectrum slicing: the
// eigenvalue count below a shift is the inertia of the banded factorization
// (banded.ts), which fixes slices of about `sliceSize` modes between zero
// and the cutoff, and each slice is solved by shift-invert Lanczos
// (lanczos.ts). Inertia counts and slices are independent jobs, run here
// through a JobRunner: in-process for the tests, or a pool of Web Workers
// in the browser (pool.ts).

import { type LinearSystem } from "../model/assemble.ts";
import { inertia } from "./banded.ts";
import { sliceModes } from "./lanczos.ts";

export type Job =
  | { kind: "inertia"; shift: number }
  | { kind: "slice"; lo: number; hi: number; countLo: number; countHi: number; seed: number };

export type JobResult =
  | { kind: "inertia"; count: number }
  | { kind: "slice"; lambda: number[]; shapes: Float64Array[]; expected: number; steps: number };

export interface JobRunner {
  run(job: Job): Promise<JobResult>;
  close(): void;
}

/** Runs the jobs on the calling thread. */
export function localRunner(sys: LinearSystem): JobRunner {
  return {
    run: async (job) => runJob(sys, job),
    close: () => {},
  };
}

export function runJob(sys: LinearSystem, job: Job): JobResult {
  if (job.kind === "inertia") return { kind: "inertia", count: inertia(sys.band, sys.mass, sys.n, sys.bw, job.shift) };
  const r = sliceModes(sys.band, sys.mass, sys.n, sys.bw, job.lo, job.hi, job.countLo, job.countHi, { seed: job.seed });
  return { kind: "slice", lambda: r.lambda, shapes: r.shapes, expected: r.expected, steps: r.steps };
}

export interface ModeSet {
  /** Number of unknowns of the system the shapes live on. */
  n: number;
  count: number;
  /** Frequencies [Hz], ascending. */
  freqs: Float64Array;
  /** Mass-normalized shapes, mode-major: shapes[m * n + i]. */
  shapes: Float32Array;
  /** Modes the inertia counts announced but Lanczos did not deliver (zero when all is well). */
  missing: number;
  /** Modes below cutoff according to the inertia count (including any missing ones). */
  expected: number;
  /** Wall time of the computation [s]. */
  seconds: number;
}

export interface ModeOptions {
  /** Target number of modes per slice; by default defaultSliceSize(bw), which balances the factorizations against the Lanczos work. */
  sliceSize?: number;
  /** Progress callback with the number of finished and total slices. */
  onProgress?: (done: number, total: number) => void;
  /** Rigid-body modes of a free plate lie at zero; modes below this frequency are discarded [Hz]. Default 1. */
  minFreq?: number;
}

/**
 * The slice size that balances the cost of a slice's factorization, about
 * n bw^2, against the Lanczos work per mode, whose reorthogonalization grows
 * with the square of the slice: the optimum is near bw / 3.5, so that a
 * finer grid, whose band is wider, uses fewer and larger slices.
 */
export function defaultSliceSize(bw: number): number {
  return Math.max(20, Math.round(bw / 3.5));
}

interface Slice {
  lo: number;
  hi: number;
  countLo: number;
  countHi: number;
}

/**
 * All modes of the system with frequencies up to fMax. The zero-frequency
 * rigid modes of a free plate are computed and then discarded.
 */
export async function computeModes(sys: LinearSystem, fMax: number, runner: JobRunner, opts: ModeOptions = {}): Promise<ModeSet> {
  const t0 = performance.now();
  const sliceSize = opts.sliceSize ?? defaultSliceSize(sys.bw);
  const minFreq = opts.minFreq ?? 1;
  const w2 = (f: number) => (2 * Math.PI * f) ** 2;
  const lamMin = -w2(1); // slightly below zero so that the rigid modes are counted in the first slice
  const lamMax = w2(fMax);
  const count = async (shift: number) => {
    const r = await runner.run({ kind: "inertia", shift });
    return (r as { count: number }).count;
  };
  const [nMin, nMax] = await Promise.all([count(lamMin), count(lamMax)]);
  const expected = nMax - nMin;
  const empty = (): ModeSet => ({ n: sys.n, count: 0, freqs: new Float64Array(0), shapes: new Float32Array(0), missing: 0, expected, seconds: (performance.now() - t0) / 1000 });
  if (expected <= 0) return empty();

  // Boundaries equally spaced in omega (the mode count of a plate grows about
  // linearly with frequency), refined where a slice is too full.
  let slices: Slice[] = [];
  {
    const S = Math.max(1, Math.ceil(expected / sliceSize));
    const wMax = Math.sqrt(lamMax);
    const bounds: number[] = [lamMin];
    for (let i = 1; i < S; i++) bounds.push(((wMax * i) / S) ** 2);
    bounds.push(lamMax);
    const counts = [nMin, ...(await Promise.all(bounds.slice(1, -1).map(count))), nMax];
    for (let i = 0; i < bounds.length - 1; i++) slices.push({ lo: bounds[i], hi: bounds[i + 1], countLo: counts[i], countHi: counts[i + 1] });
  }
  for (let round = 0; round < 6; round++) {
    const full = slices.filter((s) => s.countHi - s.countLo > 1.5 * sliceSize && s.hi - s.lo > w2(0.5));
    if (full.length === 0) break;
    const mids = await Promise.all(full.map((s) => count(((Math.sqrt(Math.max(s.lo, 0)) + Math.sqrt(s.hi)) / 2) ** 2)));
    const next: Slice[] = [];
    for (const s of slices) {
      const idx = full.indexOf(s);
      if (idx < 0) { next.push(s); continue; }
      const mid = ((Math.sqrt(Math.max(s.lo, 0)) + Math.sqrt(s.hi)) / 2) ** 2;
      next.push({ lo: s.lo, hi: mid, countLo: s.countLo, countHi: mids[idx] }, { lo: mid, hi: s.hi, countLo: mids[idx], countHi: s.countHi });
    }
    slices = next;
  }
  // Merge neighbors while the union stays within the target, and drop empty slices.
  {
    const merged: Slice[] = [];
    for (const s of slices) {
      const last = merged[merged.length - 1];
      if (last && last.countHi === s.countLo && s.countHi - last.countLo <= sliceSize) { last.hi = s.hi; last.countHi = s.countHi; }
      else merged.push({ ...s });
    }
    slices = merged.filter((s) => s.countHi > s.countLo);
  }

  let done = 0;
  opts.onProgress?.(0, slices.length);
  const results = await Promise.all(
    slices.map(async (s, i) => {
      const r = (await runner.run({ kind: "slice", lo: s.lo, hi: s.hi, countLo: s.countLo, countHi: s.countHi, seed: 1000 + 7919 * i })) as Extract<JobResult, { kind: "slice" }>;
      done++;
      opts.onProgress?.(done, slices.length);
      return r;
    }),
  );
  const pairs: { lambda: number; shape: Float64Array }[] = [];
  let missing = 0;
  for (const r of results) {
    missing += r.expected - r.lambda.length;
    for (let i = 0; i < r.lambda.length; i++) pairs.push({ lambda: r.lambda[i], shape: r.shapes[i] });
  }
  pairs.sort((a, b) => a.lambda - b.lambda);
  const kept = pairs.filter((p) => p.lambda > w2(minFreq));
  const n = sys.n;
  const shapes = new Float32Array(kept.length * n);
  const freqs = new Float64Array(kept.length);
  kept.forEach((p, m) => {
    freqs[m] = Math.sqrt(p.lambda) / (2 * Math.PI);
    shapes.set(p.shape, m * n);
  });
  return { n, count: kept.length, freqs, shapes, missing, expected, seconds: (performance.now() - t0) / 1000 };
}