// Time-domain RK4 integration of the coupled mallet / body / cavity system.
// Port of simulate() in woodblock.py. State: [q(N), qdot(N), xi, xidot, y, ydot].

import type { WoodblockParams } from '../params';
import type { ModalBasisData } from './basis';

export interface SolveResult {
  Qdd: Float64Array[];   // N arrays of length T: modal accelerations
  Xdd: Float64Array;     // slit-plug acceleration
  force: Float64Array;   // contact force f_c(t)
  fsHi: number;
  helmholtzHz: number;
  peakForceN: number;
  contactMs: number;
}

export function simulate(
  cfg: WoodblockParams,
  basis: ModalBasisData,
  oversample: number,
  onProgress?: (frac: number) => void,
): SolveResult {
  const N = basis.n;
  if (N === 0) throw new Error('No modes below max_mode_hz -- check geometry/material.');

  const fsHi = cfg.sample_rate * oversample;
  const dt = 1.0 / fsHi;
  const nSteps = Math.floor(cfg.duration_s * fsHi);

  const om2 = new Float64Array(N);
  const twoZetaOm = new Float64Array(N);
  for (let i = 0; i < N; i++) {
    om2[i] = basis.omega[i] * basis.omega[i];
    twoZetaOm[i] = 2.0 * basis.zeta[i] * basis.omega[i];
  }
  const phiC = basis.phiC;
  const A = basis.A;

  const rho0 = cfg.air_density;
  const c = cfg.speed_of_sound;
  const V = cfg.cavity_volume_m3;
  const S = cfg.slit_area_m2;
  // End correction on the slit: ~0.85 * equivalent radius per open end.
  const Leff = cfg.slit_length_m + 2.0 * 0.85 * Math.sqrt(S / Math.PI);
  const omegaH = c * Math.sqrt(S / (V * Leff));
  const twoZetaHOm = 2.0 * cfg.cavity_damping * omegaH;
  const omegaH2 = omegaH * omegaH;
  const pScale = (rho0 * c * c) / V;
  const cavDrive = (c * c) / (V * Leff);

  const K = cfg.mallet_stiffness;
  const alpha = cfg.contact_exponent;
  const mu = cfg.contact_damping;
  const invMh = 1.0 / cfg.mallet_mass;

  const dim = 2 * N + 4;
  const IX = 2 * N;
  const IXD = 2 * N + 1;
  const IY = 2 * N + 2;
  const IYD = 2 * N + 3;

  const state = new Float64Array(dim);
  state[IYD] = cfg.mallet_velocity;

  let lastFc = 0.0;
  let lastXidd = 0.0;

  const qddScratch = new Float64Array(N);

  function deriv(s: Float64Array, d: Float64Array): void {
    // Mallet contact
    let w = 0.0;
    let wd = 0.0;
    for (let i = 0; i < N; i++) {
      w += phiC[i] * s[i];
      wd += phiC[i] * s[N + i];
    }
    const y = s[IY];
    const yd = s[IYD];
    const delta = y - w;
    let fc = 0.0;
    if (delta > 0.0) {
      // Hunt-Crossley: f = K d^a (1 + lambda * ddot); lambda is in s/m so the
      // correction stays dimensionless and O(1) over a sub-ms collision.
      fc = K * Math.pow(delta, alpha) * (1.0 + mu * (yd - wd));
      if (fc < 0.0) fc = 0.0;
    }
    lastFc = fc;

    // Cavity pressure from combined volume change of walls and slit plug
    let Aq = 0.0;
    for (let i = 0; i < N; i++) Aq += A[i] * s[i];
    const pc = pScale * (S * s[IX] + Aq);

    for (let i = 0; i < N; i++) {
      const qdd = -twoZetaOm[i] * s[N + i] - om2[i] * s[i] + fc * phiC[i] - A[i] * pc;
      d[i] = s[N + i];
      d[N + i] = qdd;
      qddScratch[i] = qdd;
    }
    const xidd = -twoZetaHOm * s[IXD] - omegaH2 * s[IX] - cavDrive * Aq;
    d[IX] = s[IXD];
    d[IXD] = xidd;
    d[IY] = yd;
    d[IYD] = -fc * invMh;
    lastXidd = xidd;
  }

  const Qdd: Float64Array[] = [];
  for (let i = 0; i < N; i++) Qdd.push(new Float64Array(nSteps));
  const Xdd = new Float64Array(nSteps);
  const force = new Float64Array(nSteps);

  const k1 = new Float64Array(dim);
  const k2 = new Float64Array(dim);
  const k3 = new Float64Array(dim);
  const k4 = new Float64Array(dim);
  const tmp = new Float64Array(dim);

  const progressEvery = Math.max(1, Math.floor(nSteps / 50));

  for (let k = 0; k < nSteps; k++) {
    deriv(state, k1);
    for (let i = 0; i < N; i++) Qdd[i][k] = qddScratch[i];
    Xdd[k] = lastXidd;
    force[k] = lastFc;

    for (let i = 0; i < dim; i++) tmp[i] = state[i] + 0.5 * dt * k1[i];
    deriv(tmp, k2);
    for (let i = 0; i < dim; i++) tmp[i] = state[i] + 0.5 * dt * k2[i];
    deriv(tmp, k3);
    for (let i = 0; i < dim; i++) tmp[i] = state[i] + dt * k3[i];
    deriv(tmp, k4);
    for (let i = 0; i < dim; i++) {
      state[i] += (dt / 6.0) * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i]);
    }

    if (onProgress && k % progressEvery === 0) onProgress(k / nSteps);
  }

  let peakForce = 0.0;
  let contactSamples = 0;
  for (let k = 0; k < nSteps; k++) {
    if (force[k] > peakForce) peakForce = force[k];
    if (force[k] > 0.0) contactSamples++;
  }

  return {
    Qdd,
    Xdd,
    force,
    fsHi,
    helmholtzHz: omegaH / (2 * Math.PI),
    peakForceN: peakForce,
    contactMs: (contactSamples * dt) * 1e3,
  };
}