import type { BeadSystem, Excitation, Mallet, SimOptions, SimResult } from './types';

// Upper bound on the highest angular frequency of the linearized system
// (Gershgorin on M^-1 K at rest). A spring under tension T has stiffness k along
// its axis and T/l across it, so each spring contributes 2 max(k, |T|/l) / m.
export function omegaMaxBound(sys: BeadSystem, mallet?: Mallet): number {
  const { n, si, sj, sk, sL0 } = sys;
  const rowSum = new Float64Array(n);
  const p = (node: number, c: number) =>
    node < n ? sys.rest[3 * node + c] : sys.anchors[node - n].pos[c];
  for (let s = 0; s < si.length; s++) {
    const a = si[s];
    const b = sj[s];
    let l2 = 0;
    for (let c = 0; c < 3; c++) l2 += (p(b, c) - p(a, c)) ** 2;
    const l = Math.sqrt(l2);
    const T = sk[s] * (l - sL0[s]);
    const kmax = Math.max(sk[s], Math.abs(T) / l);
    if (a < n) rowSum[a] += 2 * kmax;
    if (b < n) rowSum[b] += 2 * kmax;
  }
  // a mallet adds a rank-one contact stiffness k w w^T between it and the beads,
  // linearized at the largest compression a rigid surface would see
  let kc = 0;
  if (mallet) {
    const { K, p, mass, v0 } = mallet;
    const dmax = (((p + 1) * mass * v0 * v0) / (2 * K)) ** (1 / (p + 1));
    kc = p * K * dmax ** (p - 1);
  }
  let w2 = mallet ? kc / mallet.mass : 0;
  for (let i = 0; i < n; i++) {
    if (sys.pinned[i]) continue;
    const extra = mallet ? kc * mallet.weights[i] : 0;
    w2 = Math.max(w2, (rowSum[i] + extra) / sys.mass[i]);
  }
  return Math.sqrt(w2);
}

export type Progress = (fraction: number) => void;

// Velocity Verlet (kick-drift-kick), explicit. Spring forces are exact
// (geometric nonlinearity kept); damping uses the half-step velocity in the
// kick, and the full-step velocity whenever an audio sample is recorded.
export function simulate(
  sys: BeadSystem,
  exc: Excitation,
  opt: SimOptions,
  progress?: Progress,
): SimResult {
  const t0 = performance.now();
  const { n, si, sj, sk, sL0, sc, mass, pinned, anchors } = sys;
  const nn = n + anchors.length;
  const ns = si.length;

  const mallet = exc.mallet;
  const omegaMax = omegaMaxBound(sys, mallet);
  const dtMax = (opt.cfl * 2) / omegaMax;
  const stepsPerSample = Math.max(1, Math.ceil(1 / (opt.fs * dtMax)));
  const dt = 1 / (opt.fs * stepsPerSample);
  const nSamples = Math.round(opt.duration * opt.fs);
  const samplesPerFrame = Math.max(1, Math.round(opt.fs / opt.frameRate));
  const nFrames = Math.floor((nSamples - 1) / samplesPerFrame) + 1;

  const X = new Float64Array(3 * nn);
  const V = new Float64Array(3 * nn);
  const F = new Float64Array(3 * n); // force on beads
  const A = new Float64Array(3 * n);
  for (let i = 0; i < n; i++) {
    for (let c = 0; c < 3; c++) {
      const k = 3 * i + c;
      X[k] = sys.rest[k] + (pinned[i] ? 0 : exc.u0[k]);
      V[k] = pinned[i] ? 0 : exc.v0[k];
    }
  }
  const slide: number[] = [];
  for (let a = 0; a < anchors.length; a++) {
    for (let c = 0; c < 3; c++) X[3 * (n + a) + c] = anchors[a].pos[c];
    if (anchors[a].slideBead >= 0) slide.push(a);
  }
  const syncAnchors = () => {
    for (const a of slide) {
      const b = anchors[a].slideBead;
      const k = 3 * (n + a);
      X[k + 1] = X[3 * b + 1];
      X[k + 2] = X[3 * b + 2];
      V[k + 1] = V[3 * b + 1];
      V[k + 2] = V[3 * b + 2];
    }
  };
  syncAnchors();

  const force = exc.force;
  const fShape = force?.shape;
  const pulse = (t: number) =>
    force && t >= 0 && t <= force.tau
      ? force.F0 * 0.5 * (1 - Math.cos((2 * Math.PI * t) / force.tau))
      : 0;

  // mallet state: position s and velocity along dir; contact force and its potential
  let ms = 0;
  let mv = mallet ? mallet.v0 : 0;
  let ma = 0;
  let fContact = 0;
  let eContact = 0;
  const mw = mallet?.weights;
  const [mdx, mdy, mdz] = mallet?.dir ?? [0, 0, 0];
  const contact = () => {
    fContact = 0;
    eContact = 0;
    if (!mallet) return;
    let h = 0;
    for (let i = 0; i < n; i++) {
      const wi = mw![i];
      if (wi === 0) continue;
      const k = 3 * i;
      h += wi * ((X[k] - sys.rest[k]) * mdx + (X[k + 1] - sys.rest[k + 1]) * mdy + (X[k + 2] - sys.rest[k + 2]) * mdz);
    }
    const d = ms - h;
    if (d > 0) {
      const dp = mallet.K * d ** mallet.p;
      fContact = dp;
      eContact = (dp * d) / (mallet.p + 1);
      for (let i = 0; i < n; i++) {
        const wi = mw![i];
        if (wi === 0) continue;
        F[3 * i] += fContact * wi * mdx;
        F[3 * i + 1] += fContact * wi * mdy;
        F[3 * i + 2] += fContact * wi * mdz;
      }
    }
    ma = -fContact / mallet.mass;
  };

  // Fills F with total force on beads; returns damping power (W, >= 0)
  // and external power (W), both at velocity V.
  let pDamp = 0;
  let pExt = 0;
  const computeForces = (t: number) => {
    F.fill(0);
    contact();
    pDamp = 0;
    for (let s = 0; s < ns; s++) {
      const a = 3 * si[s];
      const b = 3 * sj[s];
      const dx = X[b] - X[a];
      const dy = X[b + 1] - X[a + 1];
      const dz = X[b + 2] - X[a + 2];
      const l = Math.sqrt(dx * dx + dy * dy + dz * dz);
      const g = (sk[s] * (l - sL0[s])) / l;
      let fx = g * dx;
      let fy = g * dy;
      let fz = g * dz;
      const c = sc[s];
      if (c !== 0) {
        const rx = V[b] - V[a];
        const ry = V[b + 1] - V[a + 1];
        const rz = V[b + 2] - V[a + 2];
        fx += c * rx;
        fy += c * ry;
        fz += c * rz;
        pDamp += c * (rx * rx + ry * ry + rz * rz);
      }
      if (si[s] < n) {
        F[a] += fx;
        F[a + 1] += fy;
        F[a + 2] += fz;
      }
      if (sj[s] < n) {
        F[b] -= fx;
        F[b + 1] -= fy;
        F[b + 2] -= fz;
      }
    }
    const s0 = sys.sigma0;
    const p = fShape ? pulse(t) : 0;
    pExt = 0;
    for (let i = 0; i < n; i++) {
      const k = 3 * i;
      if (pinned[i]) {
        F[k] = F[k + 1] = F[k + 2] = 0;
        continue;
      }
      if (s0 !== 0) {
        const m2 = 2 * s0 * mass[i];
        for (let c = 0; c < 3; c++) {
          F[k + c] -= m2 * V[k + c];
          pDamp += m2 * V[k + c] * V[k + c];
        }
      }
      if (p !== 0) {
        for (let c = 0; c < 3; c++) {
          const fe = p * fShape![k + c];
          F[k + c] += fe;
          pExt += fe * V[k + c];
        }
      }
    }
    for (let i = 0; i < n; i++) {
      const im = 1 / mass[i];
      A[3 * i] = F[3 * i] * im;
      A[3 * i + 1] = F[3 * i + 1] * im;
      A[3 * i + 2] = F[3 * i + 2] * im;
    }
  };

  const damped = sys.sigma0 !== 0 || sc.some((c) => c !== 0);
  let maxStretch = 0;
  const energyNow = () => {
    let e = 0;
    for (let i = 0; i < n; i++) {
      const k = 3 * i;
      e += 0.5 * mass[i] * (V[k] * V[k] + V[k + 1] * V[k + 1] + V[k + 2] * V[k + 2]);
    }
    for (let s = 0; s < ns; s++) {
      const a = 3 * si[s];
      const b = 3 * sj[s];
      const l = Math.hypot(X[b] - X[a], X[b + 1] - X[a + 1], X[b + 2] - X[a + 2]);
      const d = l - sL0[s];
      e += 0.5 * sk[s] * d * d;
      maxStretch = Math.max(maxStretch, Math.abs(d) / sL0[s]);
    }
    if (mallet) e += 0.5 * mallet.mass * mv * mv + eContact;
    return e;
  };

  const audio = new Float32Array(nSamples);
  const frames = new Float32Array(nFrames * 3 * n);
  const frameTimes = new Float64Array(nFrames);
  const energy = new Float64Array(nFrames);
  const balance = new Float64Array(nFrames);
  const contactForce = mallet ? new Float32Array(nSamples) : undefined;
  const malletPos = mallet ? new Float32Array(nFrames) : undefined;
  const w = opt.pickup;

  // energy is reported relative to the equilibrium state (pre-tension energy subtracted)
  let Eeq = 0;
  for (let s = 0; s < ns; s++) {
    const pa = (node: number, c: number) => (node < n ? sys.rest[3 * node + c] : anchors[node - n].pos[c]);
    const l = Math.hypot(pa(sj[s], 0) - pa(si[s], 0), pa(sj[s], 1) - pa(si[s], 1), pa(sj[s], 2) - pa(si[s], 2));
    Eeq += 0.5 * sk[s] * (l - sL0[s]) ** 2;
  }
  computeForces(0); // contact energy at t = 0
  const E0 = energyNow();
  let dissipated = 0;
  let workIn = 0;
  let t = 0;
  computeForces(0);
  let frame = 0;
  const progressEvery = Math.max(1, Math.floor(nSamples / 100));

  for (let smp = 0; smp < nSamples; smp++) {
    // record at t = smp / fs; A holds forces at full-step velocity here
    let acc = 0;
    for (let k = 0; k < 3 * n; k++) acc += w[k] * A[k];
    audio[smp] = acc;
    if (contactForce) contactForce[smp] = fContact;
    if (smp % samplesPerFrame === 0 && frame < nFrames) {
      if (malletPos) malletPos[frame] = ms;
      const off = frame * 3 * n;
      for (let k = 0; k < 3 * n; k++) frames[off + k] = X[k] - sys.rest[k];
      frameTimes[frame] = t;
      const eTotal = energyNow();
      // the plot shows the instrument's own energy; the balance includes the mallet
      energy[frame] = eTotal - Eeq - (mallet ? 0.5 * mallet.mass * mv * mv + eContact : 0);
      balance[frame] = eTotal + dissipated - workIn - E0;
      frame++;
    }
    if (smp === nSamples - 1) break;

    for (let st = 0; st < stepsPerSample; st++) {
      // kick (half), drift
      for (let k = 0; k < 3 * n; k++) {
        V[k] += 0.5 * dt * A[k];
        X[k] += dt * V[k];
      }
      if (mallet) {
        mv += 0.5 * dt * ma;
        ms += dt * mv;
      }
      syncAnchors();
      t += dt;
      // forces at new position with half-step velocity; powers use it too
      computeForces(t);
      dissipated += dt * pDamp;
      workIn += dt * pExt;
      for (let k = 0; k < 3 * n; k++) V[k] += 0.5 * dt * A[k];
      if (mallet) mv += 0.5 * dt * ma;
      syncAnchors();
    }
    // at a sample time: recompute with the full-step velocity
    if (damped) computeForces(t);
    if (progress && smp % progressEvery === 0) progress(smp / nSamples);
  }

  return {
    audio,
    fs: opt.fs,
    frames,
    frameTimes,
    energy,
    balance,
    dt,
    stepsPerSample,
    omegaMax,
    wallSeconds: (performance.now() - t0) / 1000,
    maxStretch,
    contactForce,
    malletPos,
  };
}