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

// Normal speed of a mallet into the surface, and its linearized contact
// stiffness at the largest compression a rigid surface would see.
function contactStiffness(m: Mallet): number {
  const vn = m.vel0[0] * m.normal[0] + m.vel0[1] * m.normal[1] + m.vel0[2] * m.normal[2];
  if (!(vn > 0)) return 0;
  const dmax = (((m.p + 1) * m.mass * vn * vn) / (2 * m.K)) ** (1 / (m.p + 1));
  return m.p * m.K * dmax ** (m.p - 1);
}

function nodePositions(sys: BeadSystem): Float64Array {
  const X = new Float64Array(3 * (sys.n + sys.anchors.length));
  X.set(sys.rest);
  sys.anchors.forEach((a, k) => X.set(a.pos, 3 * (sys.n + k)));
  return X;
}

// 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.
// A hinge contributes k g g^T (g = gradient of its angle); a mallet a rank-one
// contact stiffness between it and the patch.
export function omegaMaxBound(sys: BeadSystem, mallet?: Mallet): number {
  const { n, si, sj, sk, sL0 } = sys;
  const rowSum = new Float64Array(n);
  const X = nodePositions(sys);
  for (let s = 0; s < si.length; s++) {
    const a = si[s];
    const b = sj[s];
    const l = Math.hypot(X[3 * b] - X[3 * a], X[3 * b + 1] - X[3 * a + 1], X[3 * b + 2] - X[3 * a + 2]);
    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;
  }
  const h = sys.bending;
  if (h) {
    const kh = hingeStiffness(h);
    const g = new Float64Array(13);
    for (let e = 0; e < kh.length; e++) {
      const nodes = [h.w0[e], h.w1[e], h.e0[e], h.e1[e]];
      hingeAngleGrad(X, nodes[0], nodes[1], nodes[2], nodes[3], g);
      let l1 = 0;
      for (let q = 1; q < 13; q++) l1 += Math.abs(g[q]);
      nodes.forEach((node, k) => {
        const inf = Math.max(Math.abs(g[1 + 3 * k]), Math.abs(g[2 + 3 * k]), Math.abs(g[3 + 3 * k]));
        if (node < n) rowSum[node] += kh[e] * inf * l1;
      });
    }
  }
  const kc = mallet ? contactStiffness(mallet) : 0;
  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);
}

// per hinge: sum of |M| entries of its rows over the triangles that hold it
function hingeStiffness(b: Bending): Float64Array {
  const k = new Float64Array(b.w0.length);
  for (let t = 0; t < b.tri.length / 3; t++)
    for (let a = 0; a < 3; a++) {
      const e = b.tri[3 * t + a];
      if (e < 0) continue;
      for (let c = 0; c < 3; c++) k[e] += Math.abs(b.M[9 * t + 3 * a + c]);
    }
  return k;
}

// bending energy at positions X, optionally only over triangles with keep(t)
export function bendingEnergy(b: Bending, X: Float64Array, keep?: (t: number) => boolean): number {
  const g = new Float64Array(13);
  const phi = Float64Array.from(b.w0, (_, e) => b.sign[e] * (hingeAngleGrad(X, b.w0[e], b.w1[e], b.e0[e], b.e1[e], g) - b.theta0[e]));
  let E = 0;
  for (let t = 0; t < b.tri.length / 3; t++) {
    if (keep && !keep(t)) continue;
    for (let a = 0; a < 3; a++)
      for (let c = 0; c < 3; c++) {
        const ea = b.tri[3 * t + a];
        const ec = b.tri[3 * t + c];
        if (ea >= 0 && ec >= 0) E += 0.5 * phi[ea] * b.M[9 * t + 3 * a + c] * phi[ec];
      }
  }
  return E;
}

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

// Velocity Verlet (kick-drift-kick), explicit. Spring and hinge forces are
// exact (geometric nonlinearity kept); damping and friction use 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 bend = sys.bending;
  const nh = bend ? bend.w0.length : 0;
  const nt = bend ? bend.tri.length / 3 : 0;
  const hG = new Float64Array(12 * nh); // gradients of phi
  const hPhi = new Float64Array(nh);
  const hRate = new Float64Array(nh);
  const hMom = new Float64Array(nh);

  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: 3D position q, velocity qv, acceleration qa
  const q = new Float64Array(3);
  const qv = new Float64Array(mallet ? mallet.vel0 : [0, 0, 0]);
  const qa = new Float64Array(3);
  let fContact = 0;
  let eContact = 0;
  let pFriction = 0;
  const patch: number[] = [];
  let mEff = 1;
  if (mallet) {
    let inv = 0;
    for (let i = 0; i < n; i++)
      if (mallet.weights[i] > 0) {
        patch.push(i);
        inv += (mallet.weights[i] * mallet.weights[i]) / mass[i];
      }
    mEff = 1 / (1 / mallet.mass + inv);
  }
  const [nx, ny, nz] = mallet?.normal ?? [0, 0, 0];
  const contact = () => {
    fContact = 0;
    eContact = 0;
    pFriction = 0;
    qa.fill(0);
    if (!mallet) return;
    const w = mallet.weights;
    let ux = 0, uy = 0, uz = 0, vx = 0, vy = 0, vz = 0;
    for (const i of patch) {
      const k = 3 * i;
      ux += w[i] * (X[k] - sys.rest[k]);
      uy += w[i] * (X[k + 1] - sys.rest[k + 1]);
      uz += w[i] * (X[k + 2] - sys.rest[k + 2]);
      vx += w[i] * V[k];
      vy += w[i] * V[k + 1];
      vz += w[i] * V[k + 2];
    }
    const d = (q[0] - ux) * nx + (q[1] - uy) * ny + (q[2] - uz) * nz;
    if (!(d > 0)) return;
    if (mallet.reach !== undefined) {
      const ox = q[0] - ux - d * nx, oy = q[1] - uy - d * ny, oz = q[2] - uz - d * nz;
      if (ox * ox + oy * oy + oz * oz > mallet.reach * mallet.reach) return;
    }
    const fn = mallet.K * d ** mallet.p;
    fContact = fn;
    eContact = (fn * d) / (mallet.p + 1);
    // force on the mallet: -fn normal + friction against tangential slip
    let fx = -fn * nx;
    let fy = -fn * ny;
    let fz = -fn * nz;
    if (mallet.mu > 0) {
      let rx = qv[0] - vx, ry = qv[1] - vy, rz = qv[2] - vz;
      const rn = rx * nx + ry * ny + rz * nz;
      rx -= rn * nx;
      ry -= rn * ny;
      rz -= rn * nz;
      const slip = Math.hypot(rx, ry, rz);
      // regularized Coulomb; the slip scale keeps the explicit step stable
      const eps = Math.max(1e-3, (4 * mallet.mu * fn * dt) / mEff);
      const g = (mallet.mu * fn) / Math.sqrt(slip * slip + eps * eps);
      fx -= g * rx;
      fy -= g * ry;
      fz -= g * rz;
      pFriction = g * slip * slip;
    }
    qa[0] = fx / mallet.mass;
    qa[1] = fy / mallet.mass;
    qa[2] = fz / mallet.mass;
    for (const i of patch) {
      F[3 * i] -= w[i] * fx;
      F[3 * i + 1] -= w[i] * fy;
      F[3 * i + 2] -= w[i] * fz;
    }
  };

  const hOut = new Float64Array(13);
  let eHinge = 0;
  // Fills F with total force on beads; sets 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 = pFriction;
    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;
      }
    }
    eHinge = 0;
    if (bend) {
      const s1 = bend.sigma1;
      for (let e = 0; e < nh; e++) {
        const th = hingeAngleGrad(X, bend.w0[e], bend.w1[e], bend.e0[e], bend.e1[e], hOut);
        const sg = bend.sign[e];
        hPhi[e] = sg * (th - bend.theta0[e]);
        const o = 12 * e;
        for (let q = 0; q < 12; q++) hG[o + q] = sg * hOut[1 + q];
        if (s1 !== 0) {
          const w0 = 3 * bend.w0[e], w1 = 3 * bend.w1[e], e0 = 3 * bend.e0[e], e1 = 3 * bend.e1[e];
          hRate[e] =
            hG[o] * V[w0] + hG[o + 1] * V[w0 + 1] + hG[o + 2] * V[w0 + 2] +
            hG[o + 3] * V[w1] + hG[o + 4] * V[w1 + 1] + hG[o + 5] * V[w1 + 2] +
            hG[o + 6] * V[e0] + hG[o + 7] * V[e0 + 1] + hG[o + 8] * V[e0 + 2] +
            hG[o + 9] * V[e1] + hG[o + 10] * V[e1 + 1] + hG[o + 11] * V[e1 + 2];
        }
      }
      hMom.fill(0);
      const tri = bend.tri;
      const Mt = bend.M;
      for (let t = 0; t < nt; t++) {
        for (let a = 0; a < 3; a++) {
          const ea = tri[3 * t + a];
          if (ea < 0) continue;
          let m = 0;
          let r = 0;
          for (let c = 0; c < 3; c++) {
            const ec = tri[3 * t + c];
            if (ec < 0) continue;
            const M = Mt[9 * t + 3 * a + c];
            m += M * hPhi[ec];
            r += M * hRate[ec];
          }
          eHinge += 0.5 * hPhi[ea] * m;
          if (s1 !== 0) {
            pDamp += 2 * s1 * hRate[ea] * r;
            m += 2 * s1 * r;
          }
          hMom[ea] += m;
        }
      }
      for (let e = 0; e < nh; e++) {
        const m = hMom[e];
        if (m === 0) continue;
        const o = 12 * e;
        const w0 = 3 * bend.w0[e], w1 = 3 * bend.w1[e], e0 = 3 * bend.e0[e], e1 = 3 * bend.e1[e];
        F[w0] -= m * hG[o];
        F[w0 + 1] -= m * hG[o + 1];
        F[w0 + 2] -= m * hG[o + 2];
        F[w1] -= m * hG[o + 3];
        F[w1 + 1] -= m * hG[o + 4];
        F[w1 + 2] -= m * hG[o + 5];
        F[e0] -= m * hG[o + 6];
        F[e0 + 1] -= m * hG[o + 7];
        F[e0 + 2] -= m * hG[o + 8];
        F[e1] -= m * hG[o + 9];
        F[e1 + 1] -= m * hG[o + 10];
        F[e1 + 2] -= m * hG[o + 11];
      }
    }
    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 velocityForces =
    sys.sigma0 !== 0 || sc.some((c) => c !== 0) || (bend?.sigma1 ?? 0) !== 0 || (mallet?.mu ?? 0) > 0;
  let maxStretch = 0;
  const malletEnergy = () => (mallet ? 0.5 * mallet.mass * (qv[0] ** 2 + qv[1] ** 2 + qv[2] ** 2) + eContact : 0);
  // total mechanical energy at the current state; computeForces must have run
  // at this state (it sets the hinge and contact energies)
  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;
      if (sj[s] < n && si[s] < n) maxStretch = Math.max(maxStretch, Math.abs(d) / sL0[s]);
    }
    return e + eHinge + malletEnergy();
  };

  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(3 * nFrames) : undefined;
  const w = opt.pickup;
  const delay = opt.pickupDelay;

  // 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);
  const E0 = energyNow();
  let dissipated = 0;
  let workIn = 0;
  let t = 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
    if (delay) {
      for (let i = 0; i < n; i++) {
        const k = 3 * i;
        const c = w[k] * A[k] + w[k + 1] * A[k + 1] + w[k + 2] * A[k + 2];
        if (c === 0) continue;
        const pos = smp + delay[i];
        const i0 = Math.floor(pos);
        const fr = pos - i0;
        if (i0 < nSamples) audio[i0] += (1 - fr) * c;
        if (i0 + 1 < nSamples) audio[i0 + 1] += fr * c;
      }
    } else {
      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.set(q, 3 * frame);
      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 - malletEnergy();
      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)
        for (let c = 0; c < 3; c++) {
          qv[c] += 0.5 * dt * qa[c];
          q[c] += dt * qv[c];
        }
      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) for (let c = 0; c < 3; c++) qv[c] += 0.5 * dt * qa[c];
      syncAnchors();
    }
    // at a sample time: recompute with the full-step velocity
    if (velocityForces) 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,
  };
}