// Collects the modal snapshots of a run and turns them into per-mode energies.
//
// A snapshot holds q_m and qdot_m for every mode (projections of w and w-dot onto the
// mass-weighted shapes). The energy of a mode is E_m = (qdot_m^2 + omega_m^2 q_m^2) / 2, exact
// for the linearised, condensed model; in the nonlinear run it is the standard measure of how
// much of the transverse motion lives in mode m. The sum over modes misses the in-plane and
// rotational kinetic energy and the part of the strain energy that the linear basis does not
// represent, so it is compared against, not equated to, the simulator's energy diagnostics.
import type { ChunkResult } from "../model/simulator.ts";

export class ModalTracker {
  readonly nModes: number;
  readonly omega: Float64Array;      // (Nm,) [rad/s] at the current material
  readonly times: number[] = [];     // [s]
  readonly energies: Float32Array[] = [];   // per snapshot, (Nm,) [J]
  readonly totals: number[] = [];    // sum over modes per snapshot [J]
  maxEnergy = 0;
  drivenMode: number;

  constructor(freq: Float64Array, drivenMode: number) {
    this.nModes = freq.length;
    this.omega = Float64Array.from(freq, (f) => 2 * Math.PI * f);
    this.drivenMode = drivenMode;
  }

  addChunk(res: ChunkResult, fs: number) {
    const Nm = this.nModes;
    for (let s = 0; s < res.proj.length; s++) {
      const row = res.proj[s], E = new Float32Array(Nm);
      let total = 0;
      for (let m = 0; m < Nm; m++) {
        const q = row[m], v = row[Nm + m];
        const e = 0.5 * (v * v + this.omega[m] * this.omega[m] * q * q);
        E[m] = e; total += e;
        if (e > this.maxEnergy) this.maxEnergy = e;
      }
      this.times.push(res.projSamples[s] / fs);
      this.energies.push(E);
      this.totals.push(total);
    }
  }

  get length() { return this.times.length; }
  get duration() { return this.times.length ? this.times[this.times.length - 1] : 0; }

  /** Index of the snapshot at or before ``t``. */
  indexAtTime(t: number): number {
    const ts = this.times;
    let lo = 0, hi = ts.length - 1;
    if (hi < 0) return -1;
    while (lo < hi) { const mid = (lo + hi + 1) >> 1; if (ts[mid] <= t) lo = mid; else hi = mid - 1; }
    return lo;
  }

  /** Energy of every mode except the driven one, at snapshot ``s``. */
  otherEnergy(s: number): number {
    const E = this.energies[s];
    let sum = 0;
    for (let m = 0; m < this.nModes; m++) if (m !== this.drivenMode) sum += E[m];
    return sum;
  }

  /** The ``k`` most energetic modes at snapshot ``s``, as (mode, energy) sorted descending. */
  topModes(s: number, k: number): { mode: number; energy: number }[] {
    const E = this.energies[s];
    const idx = Array.from(E.keys()).sort((a, b) => E[b] - E[a]).slice(0, k);
    return idx.map((mode) => ({ mode, energy: E[mode] }));
  }
}