// Time integration of the bead system.
//
// The conservative part (springs, external force pulses) is integrated with
// velocity Verlet, optionally composed into Yoshida's fourth-order scheme.
// Damping (friction and dashpots) is linear in the velocities, dv/dt = -A v,
// and is applied as a separate exact step v <- exp(-tau A) v, once per audio
// sample (Strang splitting over the sample interval Delta = 1/fs). The
// splitting error is second order in Delta and scales with the damping rate,
// which is small compared with the oscillation frequencies; the conservative
// part is second order (Verlet) or fourth order (Yoshida) in the step h.

import type { BeadSystem, Method, Modes, SimResult } from "./types";

const YOSHIDA_W1 = 1 / (2 - Math.cbrt(2));
const YOSHIDA_W0 = -Math.cbrt(2) * YOSHIDA_W1;

export function stepCoefficients(method: Method): number[] {
  return method === "yoshida4" ? [YOSHIDA_W1, YOSHIDA_W0, YOSHIDA_W1] : [1];
}

// Largest h*omega for which the method is stable on x'' = -omega^2 x.
export function stabilityLimit(method: Method): number {
  if (method === "verlet") return 2;
  const cs = stepCoefficients(method);
  const stable = (h: number) => {
    // transfer matrix of one full step for omega = 1
    let m = [1, 0, 0, 1]; // [x<-x, x<-v, v<-x, v<-v]
    const mul = (a: number[], b: number[]) => [
      a[0] * b[0] + a[1] * b[2],
      a[0] * b[1] + a[1] * b[3],
      a[2] * b[0] + a[3] * b[2],
      a[2] * b[1] + a[3] * b[3],
    ];
    for (const c of cs) {
      const kick = [1, 0, (-c * h) / 2, 1];
      const drift = [1, c * h, 0, 1];
      m = mul(kick, mul(drift, mul(kick, m)));
    }
    return Math.abs(m[0] + m[3]) < 2;
  };
  let h = 1e-3;
  while (stable(h + 1e-3)) h += 1e-3;
  return h;
}

// Error of the numerical oscillation frequency for x'' = -omega^2 x, in cents
// (positive: the simulated tone is sharp).
export function frequencyErrorCents(method: Method, omega: number, h: number): number {
  let m = [1, 0, 0, 1];
  const mul = (a: number[], b: number[]) => [
    a[0] * b[0] + a[1] * b[2],
    a[0] * b[1] + a[1] * b[3],
    a[2] * b[0] + a[3] * b[2],
    a[2] * b[1] + a[3] * b[3],
  ];
  for (const c of stepCoefficients(method)) {
    const kick = [1, 0, (-c * h * omega * omega) / 2, 1];
    const drift = [1, c * h, 0, 1];
    m = mul(kick, mul(drift, mul(kick, m)));
  }
  const half = (m[0] + m[3]) / 2;
  if (Math.abs(half) >= 1) return NaN;
  const omegaNum = Math.acos(half) / h;
  return 1200 * Math.log2(omegaNum / omega);
}

export interface SimOptions {
  modes?: Modes | null;
  onProgress?: (t: number) => void;
  maxAnimBytes?: number;
}

export function simulate(sys: BeadSystem, opts: SimOptions = {}): SimResult {
  const t0 = performance.now();
  const warnings: string[] = [];
  const { n, mass, pinned } = sys;
  const N3 = 3 * n;
  const x = new Float64Array(sys.pos0);
  const v = new Float64Array(N3);
  const F = new Float64Array(N3);
  const invm = new Float64Array(n);
  for (let i = 0; i < n; i++) invm[i] = pinned[i] ? 0 : 1 / mass[i];

  // ---- forces ------------------------------------------------------------
  const { si, sj, sk, sL0, sk3 } = sys;
  const ns = si.length;
  const internalForces = (): number => {
    F.fill(0);
    let pe = 0;
    for (let s = 0; s < ns; s++) {
      const i3 = 3 * si[s];
      const j3 = 3 * sj[s];
      const dx = x[j3] - x[i3];
      const dy = x[j3 + 1] - x[i3 + 1];
      const dz = x[j3 + 2] - x[i3 + 2];
      const L = Math.sqrt(dx * dx + dy * dy + dz * dz);
      const e = L - sL0[s];
      const k = sk[s];
      const k3 = sk3[s];
      const e2 = e * e;
      pe += 0.5 * k * e2 + 0.25 * k3 * e2 * e2;
      const g = (k * e + k3 * e2 * e) / L;
      F[i3] += g * dx;
      F[i3 + 1] += g * dy;
      F[i3 + 2] += g * dz;
      F[j3] -= g * dx;
      F[j3 + 1] -= g * dy;
      F[j3 + 2] -= g * dz;
    }
    return pe;
  };
  const pulses = sys.hits.filter(h => h.duration > 0);
  const externalForces = (t: number) => {
    for (const p of pulses) {
      const u = (t - p.t) / p.duration;
      if (u <= 0 || u >= 1) continue;
      const w = 0.5 * (1 - Math.cos(2 * Math.PI * u));
      for (const b of p.beads) {
        F[3 * b] += w * p.force[0];
        F[3 * b + 1] += w * p.force[1];
        F[3 * b + 2] += w * p.force[2];
      }
    }
  };

  // ---- rest state checks ---------------------------------------------------
  let pe = internalForces();
  const peRest = pe; // energy stored by pre-tension; energies are reported relative to rest
  let residualForce = 0;
  for (let i = 0; i < n; i++) {
    if (pinned[i]) continue;
    const f = Math.hypot(F[3 * i], F[3 * i + 1], F[3 * i + 2]);
    if (f > residualForce) residualForce = f;
  }
  let maxTension = 0;
  for (let s = 0; s < ns; s++) {
    const i3 = 3 * si[s];
    const j3 = 3 * sj[s];
    const L = Math.hypot(x[j3] - x[i3], x[j3 + 1] - x[i3 + 1], x[j3 + 2] - x[i3 + 2]);
    const e = L - sL0[s];
    maxTension = Math.max(maxTension, Math.abs(sk[s] * e + sk3[s] * e * e * e));
  }
  if (residualForce > 1e-6 * Math.max(maxTension, 1e-30)) {
    warnings.push(
      `The rest positions are not an equilibrium: the largest net force on a free bead is ${residualForce.toExponential(2)} N ` +
        `(largest spring force ${maxTension.toExponential(2)} N). The beads will move before any strike.`
    );
  }

  // ---- time step -----------------------------------------------------------
  const omegaMax = estimateOmegaMax(sys);
  const fs = sys.sampleRate;
  const hLimit = stabilityLimit(sys.method) / omegaMax;
  let substeps: number;
  if (sys.dt > 0) {
    substeps = Math.max(1, Math.round(1 / (fs * sys.dt)));
    if (1 / (fs * substeps) > hLimit) {
      warnings.push(
        `dt = ${(1 / (fs * substeps)).toExponential(3)} s exceeds the stability limit ${hLimit.toExponential(3)} s.`
      );
    }
  } else {
    substeps = Math.max(1, Math.ceil(1 / (fs * sys.safety * hLimit)));
  }
  const h = 1 / (fs * substeps);
  const nSamples = Math.max(1, Math.round(sys.duration * fs));
  const nSteps = nSamples * substeps;
  const cs = stepCoefficients(sys.method);

  // ---- damping: v <- exp(-tau A) v ------------------------------------------
  const { di, dj, dc, friction } = sys;
  const nd = di.length;
  let normA = 0; // infinity-norm bound of A
  {
    const row = new Float64Array(n);
    for (let i = 0; i < n; i++) row[i] = friction[i];
    for (let d = 0; d < nd; d++) {
      row[di[d]] += 2 * dc[d];
      row[dj[d]] += 2 * dc[d];
    }
    for (let i = 0; i < n; i++) normA = Math.max(normA, row[i] * invm[i]);
  }
  const hasDamping = normA > 0;
  const frictionOnly = nd === 0;
  const term = new Float64Array(N3);
  const tmp = new Float64Array(N3);
  const fa = new Float64Array(n);
  for (let i = 0; i < n; i++) fa[i] = friction[i] * invm[i];
  const dai = new Float64Array(nd);
  const daj = new Float64Array(nd);
  for (let d = 0; d < nd; d++) {
    dai[d] = dc[d] * invm[di[d]];
    daj[d] = dc[d] * invm[dj[d]];
  }
  const applyA = (w: Float64Array, out: Float64Array) => {
    for (let i = 0; i < n; i++) {
      const a = fa[i];
      out[3 * i] = a * w[3 * i];
      out[3 * i + 1] = a * w[3 * i + 1];
      out[3 * i + 2] = a * w[3 * i + 2];
    }
    for (let d = 0; d < nd; d++) {
      const i3 = 3 * di[d];
      const j3 = 3 * dj[d];
      const ai = dai[d];
      const aj = daj[d];
      const r0 = w[i3] - w[j3];
      const r1 = w[i3 + 1] - w[j3 + 1];
      const r2 = w[i3 + 2] - w[j3 + 2];
      out[i3] += ai * r0;
      out[i3 + 1] += ai * r1;
      out[i3 + 2] += ai * r2;
      out[j3] -= aj * r0;
      out[j3 + 1] -= aj * r1;
      out[j3 + 2] -= aj * r2;
    }
  };
  const kinetic = () => {
    let ke = 0;
    for (let i = 0; i < n; i++) {
      ke += 0.5 * mass[i] * (v[3 * i] ** 2 + v[3 * i + 1] ** 2 + v[3 * i + 2] ** 2);
    }
    return ke;
  };
  let dissipated = 0;
  const applyDamping = (tau: number) => {
    if (!hasDamping || tau === 0) return;
    const keBefore = kinetic();
    if (frictionOnly) {
      for (let i = 0; i < n; i++) {
        const f = Math.exp(-tau * friction[i] * invm[i]);
        v[3 * i] *= f;
        v[3 * i + 1] *= f;
        v[3 * i + 2] *= f;
      }
    } else {
      // Taylor series with scaling; |tau_s| * normA <= 0.5 so terms shrink fast.
      const nsub = Math.max(1, Math.ceil((Math.abs(tau) * normA) / 0.5));
      const ts = tau / nsub;
      for (let r = 0; r < nsub; r++) {
        let vmax = 0;
        for (let q = 0; q < N3; q++) {
          term[q] = v[q];
          vmax = Math.max(vmax, Math.abs(v[q]));
        }
        if (vmax === 0) break;
        for (let k = 1; k < 40; k++) {
          applyA(term, tmp);
          const c = -ts / k;
          let tmax = 0;
          for (let q = 0; q < N3; q++) {
            const val = c * tmp[q];
            term[q] = val;
            v[q] += val;
            const a = Math.abs(val);
            if (a > tmax) tmax = a;
          }
          if (tmax <= 1e-17 * vmax) break;
        }
      }
    }
    dissipated += keBefore - kinetic();
  };

  // ---- recording -------------------------------------------------------------
  const audio = new Float32Array(nSamples);
  const lb = sys.listenBeads;
  const [ldx, ldy, ldz] = sys.listenDir;
  const pos0 = sys.pos0;
  const listen = () => {
    let s = 0;
    for (const b of lb) {
      s +=
        (x[3 * b] - pos0[3 * b]) * ldx +
        (x[3 * b + 1] - pos0[3 * b + 1]) * ldy +
        (x[3 * b + 2] - pos0[3 * b + 2]) * ldz;
    }
    return s / lb.length;
  };
  const frameStride = Math.max(1, Math.round(fs / sys.frameRate));
  const nFrames = Math.floor((nSamples - 1) / frameStride) + 1;
  const maxAnimBytes = opts.maxAnimBytes ?? 150e6;
  const animStride = frameStride * Math.max(1, Math.ceil((nFrames * N3 * 4) / maxAnimBytes));
  const frameTimes = new Float64Array(nFrames);
  const energy = new Float64Array(nFrames);
  const dissipatedRec = new Float64Array(nFrames);
  const animTimes: number[] = [];
  const animFrames: Float32Array[] = [];

  const modes = opts.modes ?? null;
  const nModes = modes ? modes.freqs.length : 0;
  const modeEnergy = modes ? new Float32Array(nFrames * nModes) : null;
  let bz: Float64Array | null = null;
  let bv: Float64Array | null = null;
  let om2: Float64Array | null = null;
  if (modes) {
    bz = new Float64Array(modes.free.length);
    bv = new Float64Array(modes.free.length);
    om2 = new Float64Array(nModes);
    for (let k = 0; k < nModes; k++) om2[k] = (2 * Math.PI * modes.freqs[k]) ** 2;
  }
  const recordModes = (f: number) => {
    if (!modes || !modeEnergy || !bz || !bv || !om2) return;
    const free = modes.free;
    const nf = free.length;
    for (let r = 0; r < nf; r++) {
      const b = free[r];
      const sm = Math.sqrt(mass[b]);
      bz[r] = sm * (x[3 * b + 2] - pos0[3 * b + 2]);
      bv[r] = sm * v[3 * b + 2];
    }
    const psi = modes.psi;
    for (let k = 0; k < nModes; k++) {
      let q = 0;
      let qd = 0;
      const off = k * nf;
      for (let r = 0; r < nf; r++) {
        q += psi[off + r] * bz[r];
        qd += psi[off + r] * bv[r];
      }
      modeEnergy[f * nModes + k] = 0.5 * (qd * qd + om2[k] * q * q);
    }
  };

  // ---- impulses, by audio sample ----------------------------------------------
  const impulsesAt = new Map<number, typeof sys.hits>();
  for (const hit of sys.hits) {
    if (hit.duration > 0) continue;
    const s = Math.round(hit.t * fs);
    if (s >= nSamples) continue;
    if (!impulsesAt.has(s)) impulsesAt.set(s, []);
    impulsesAt.get(s)!.push(hit);
  }
  const tExcitationEnd = Math.max(0, ...sys.hits.map(hh => hh.t + hh.duration));

  // ---- main loop ---------------------------------------------------------------
  // Each audio sample interval Delta = substeps * h is one macro step:
  // D(Delta/2) [substeps conservative steps] D(Delta/2). The damping halves of
  // neighboring macro steps are merged unless something needs the velocities
  // at the boundary (a frame record or an impulse).
  const Delta = 1 / fs;
  pe = internalForces();
  externalForces(0);
  let pendingD = 0;
  let aborted: string | null = null;
  let frame = 0;
  let lastProgress = performance.now();
  for (let s = 0; s <= nSteps; s++) {
    const t = s * h;
    const sample = s % substeps === 0 ? s / substeps : -1;
    const isFrame = sample >= 0 && sample % frameStride === 0 && frame < nFrames;
    const imps = sample >= 0 ? impulsesAt.get(sample) : undefined;
    if ((isFrame || imps) && pendingD !== 0) {
      applyDamping(pendingD);
      pendingD = 0;
    }
    if (imps) {
      for (const hit of imps) {
        for (const b of hit.beads) {
          if (pinned[b]) continue;
          v[3 * b] += hit.velocity[0];
          v[3 * b + 1] += hit.velocity[1];
          v[3 * b + 2] += hit.velocity[2];
        }
      }
    }
    if (sample >= 0 && sample < nSamples) audio[sample] = listen();
    if (isFrame) {
      frameTimes[frame] = t;
      energy[frame] = kinetic() + pe - peRest;
      dissipatedRec[frame] = dissipated;
      recordModes(frame);
      if (sample % animStride === 0) {
        const d = new Float32Array(N3);
        for (let q = 0; q < N3; q++) d[q] = x[q] - pos0[q];
        animTimes.push(t);
        animFrames.push(d);
      }
      frame++;
    }
    if (s === nSteps) break;
    if (sample >= 0) {
      pendingD += Delta / 2;
      applyDamping(pendingD);
      pendingD = 0;
    }

    let tj = t;
    for (const c of cs) {
      const kh = (c * h) / 2;
      for (let i = 0; i < n; i++) {
        const im = invm[i] * kh;
        v[3 * i] += im * F[3 * i];
        v[3 * i + 1] += im * F[3 * i + 1];
        v[3 * i + 2] += im * F[3 * i + 2];
      }
      const dh = c * h;
      for (let q = 0; q < N3; q++) x[q] += dh * v[q];
      tj += c * h;
      pe = internalForces();
      externalForces(tj);
      for (let i = 0; i < n; i++) {
        const im = invm[i] * kh;
        v[3 * i] += im * F[3 * i];
        v[3 * i + 1] += im * F[3 * i + 1];
        v[3 * i + 2] += im * F[3 * i + 2];
      }
    }
    if ((s + 1) % substeps === 0) pendingD = Delta / 2;

    if ((s & 1023) === 0) {
      if (!Number.isFinite(pe)) {
        aborted = `The solution blew up at t = ${t.toFixed(4)} s. Reduce dt (or safety) or the strike strength.`;
        break;
      }
      const now = performance.now();
      if (opts.onProgress && now - lastProgress > 200) {
        opts.onProgress(t);
        lastProgress = now;
      }
    }
  }

  // energy drift after the last excitation
  let energyDrift = 0;
  let ref = -1;
  for (let f = 0; f < frame; f++) {
    if (frameTimes[f] < tExcitationEnd) continue;
    const total = energy[f] + dissipatedRec[f];
    if (ref < 0) ref = total;
    else if (ref > 0) energyDrift = Math.max(energyDrift, Math.abs(total - ref) / ref);
  }

  return {
    audio,
    sampleRate: fs,
    dt: h,
    substeps,
    omegaMax,
    frameTimes: frameTimes.slice(0, frame),
    energy: energy.slice(0, frame),
    dissipated: dissipatedRec.slice(0, frame),
    animTimes: Float64Array.from(animTimes),
    animFrames,
    modeEnergy: modeEnergy ? modeEnergy.slice(0, frame * nModes) : null,
    residualForce,
    maxTension,
    energyDrift,
    freqErrors: [1000, 2000, 5000, 10000, 20000]
      .filter(hz => hz < fs / 2)
      .map(hz => ({ hz, cents: frequencyErrorCents(sys.method, 2 * Math.PI * hz, h) })),
    elapsedMs: performance.now() - t0,
    aborted,
    warnings,
  };
}

// Highest frequency (rad/s) of the system linearized about the rest positions,
// by power iteration on M^{-1/2} H M^{-1/2}, with a small margin.
export function estimateOmegaMax(sys: BeadSystem): number {
  const { n, si, sj, sk, sL0, sk3, pos0, mass, pinned } = sys;
  const ns = si.length;
  // per-spring 3x3 blocks B = keff d d^T + (f/L)(I - d d^T)
  const B = new Float64Array(ns * 6);
  for (let s = 0; s < ns; s++) {
    const i3 = 3 * si[s];
    const j3 = 3 * sj[s];
    const dx = pos0[j3] - pos0[i3];
    const dy = pos0[j3 + 1] - pos0[i3 + 1];
    const dz = pos0[j3 + 2] - pos0[i3 + 2];
    const L = Math.sqrt(dx * dx + dy * dy + dz * dz);
    const e = L - sL0[s];
    const keff = sk[s] + 3 * sk3[s] * e * e;
    const g = (sk[s] * e + sk3[s] * e * e * e) / L;
    const u = [dx / L, dy / L, dz / L];
    const a = keff - g;
    B[6 * s] = a * u[0] * u[0] + g;
    B[6 * s + 1] = a * u[1] * u[1] + g;
    B[6 * s + 2] = a * u[2] * u[2] + g;
    B[6 * s + 3] = a * u[0] * u[1];
    B[6 * s + 4] = a * u[0] * u[2];
    B[6 * s + 5] = a * u[1] * u[2];
  }
  const ism = new Float64Array(n);
  for (let i = 0; i < n; i++) ism[i] = pinned[i] ? 0 : 1 / Math.sqrt(mass[i]);
  let w = new Float64Array(3 * n);
  let seed = 12345;
  for (let q = 0; q < 3 * n; q++) {
    seed = (seed * 1103515245 + 12345) & 0x7fffffff;
    w[q] = ism[Math.floor(q / 3)] * (seed / 0x7fffffff - 0.5);
  }
  const u = new Float64Array(3 * n);
  const y = new Float64Array(3 * n);
  let lambda = 0;
  for (let it = 0; it < 500; it++) {
    let norm = 0;
    for (let q = 0; q < 3 * n; q++) norm += w[q] * w[q];
    norm = Math.sqrt(norm);
    if (norm === 0) return 1;
    for (let q = 0; q < 3 * n; q++) u[q] = (w[q] / norm) * ism[Math.floor(q / 3)];
    y.fill(0);
    for (let s = 0; s < ns; s++) {
      const i3 = 3 * si[s];
      const j3 = 3 * sj[s];
      const r0 = u[i3] - u[j3];
      const r1 = u[i3 + 1] - u[j3 + 1];
      const r2 = u[i3 + 2] - u[j3 + 2];
      const b = 6 * s;
      const f0 = B[b] * r0 + B[b + 3] * r1 + B[b + 4] * r2;
      const f1 = B[b + 3] * r0 + B[b + 1] * r1 + B[b + 5] * r2;
      const f2 = B[b + 4] * r0 + B[b + 5] * r1 + B[b + 2] * r2;
      y[i3] += f0;
      y[i3 + 1] += f1;
      y[i3 + 2] += f2;
      y[j3] -= f0;
      y[j3 + 1] -= f1;
      y[j3 + 2] -= f2;
    }
    let rq = 0;
    const wNext = new Float64Array(3 * n);
    for (let q = 0; q < 3 * n; q++) {
      wNext[q] = y[q] * ism[Math.floor(q / 3)];
      rq += (w[q] / norm) * wNext[q];
    }
    w = wNext;
    if (it > 20 && Math.abs(rq - lambda) < 1e-8 * Math.abs(rq)) {
      lambda = rq;
      break;
    }
    lambda = rq;
  }
  return Math.sqrt(Math.max(lambda, 1e-300)) * 1.02;
}