Inverse kinematics and joint trajectories

/** Pose error of `link` (point `local`) against target {p?, R?}: [Δp (3), rotation vector of R_t·Rᵀ (3)] in world axes. */
export function poseError(robot, q, link, target, local = [0, 0, 0], base) {
  const T = robot.pose(q, link, base), pw = M.T_apply(T, local), e = [];
  if (target.p) e.push(...M.sub(target.p, pw));
  if (target.R) e.push(...M.rotvec(M.mmul(target.R, M.mT(T.R))));
  return e;
}

/**
 * ik(robot, link, target, q0, opts) → {ok, q, iterations, error: {pos, rot}, problems}
 * target {p?: [x,y,z], R?: 3×3}: position only, orientation only, or both. opts: local, base, maxIter (300), tol (1e-6 m), rotTol (1e-6 rad),
 * lambda (0.1), weightRot (1), stepMax (0.3 rad or m per iteration), limits (true).
 */
export function ik(robot, link, target, q0 = null, opts = {}) {
  const { local = [0, 0, 0], base, maxIter = 300, tol = 1e-6, rotTol = 1e-6, weightRot = 1, stepMax = 0.3, limits = true } = opts;
  let lambda = opts.lambda ?? 0.1;
  const n = robot.dof, lim = robot.limits(), problems = [];
  if (!robot.ok) return { ok: false, q: q0 || new Array(n).fill(0), iterations: 0, error: { pos: NaN, rot: NaN }, problems: robot.problems.slice() };
  if (!target.p && !target.R) return { ok: false, q: q0 || new Array(n).fill(0), iterations: 0, error: { pos: NaN, rot: NaN }, problems: [M.problem('target', 'error', 'The target has neither a position nor an orientation.')] };
  let q = limits ? clampQ(q0 || new Array(n).fill(0), lim) : (q0 || new Array(n).fill(0)).slice();
  const rows = (target.p ? 3 : 0) + (target.R ? 3 : 0);
  const errNorm = (e) => { let i = 0, pe = 0, re = 0; if (target.p) { pe = Math.hypot(e[0], e[1], e[2]); i = 3; } if (target.R) re = Math.hypot(e[i], e[i + 1], e[i + 2]); return { pos: pe, rot: re }; };
  const weigh = (e) => { if (!target.R) return e; const o = e.slice(); const s = target.p ? 3 : 0; for (let k = 0; k < 3; k++) o[s + k] *= weightRot; return o; };
  let e = poseError(robot, q, link, target, local, base), er = errNorm(e), cost = er.pos ** 2 + (weightRot * er.rot) ** 2, it = 0;
  for (; it < maxIter; it++) {
    if (er.pos <= tol && er.rot <= rotTol) break;
    const Jf = robot.jacobian(q, link, local, base); // 6×n
    let J = []; if (target.p) J.push(Jf[0], Jf[1], Jf[2]); if (target.R) for (const r of [3, 4, 5]) J.push(Jf[r].map((v) => v * weightRot));
    let improved = false;
    for (let tries = 0; tries < 12 && !improved; tries++) {
      const A = M.zeros(rows, rows); for (let i = 0; i < rows; i++) for (let j = 0; j < rows; j++) { let s = 0; for (let k = 0; k < n; k++) s += J[i][k] * J[j][k]; A[i][j] = s + (i === j ? lambda * lambda : 0); }
      const y = M.solve(A, weigh(e)); if (!y) { lambda *= 4; continue; }
      let dq = new Array(n).fill(0); for (let k = 0; k < n; k++) for (let i = 0; i < rows; i++) dq[k] += J[i][k] * y[i];
      const big = Math.max(...dq.map(Math.abs)); if (big > stepMax) dq = dq.map((v) => v * stepMax / big);
      let qn = q.map((v, i) => v + dq[i]); if (limits) qn = clampQ(qn, lim);
      const en = poseError(robot, qn, link, target, local, base), ern = errNorm(en), cn = ern.pos ** 2 + (weightRot * ern.rot) ** 2;
      if (cn < cost) { q = qn; e = en; er = ern; cost = cn; lambda = Math.max(1e-4, lambda * 0.5); improved = true; } else lambda *= 3;
    }
    if (!improved) break; // stuck: a joint limit, a singularity or an unreachable target
  }
  const ok = er.pos <= tol * 10 && er.rot <= rotTol * 10;
  if (!ok) {
    const atLimit = limits ? robot.names.filter((_, i) => Math.abs(q[i] - lim[i].lower) < 1e-9 || Math.abs(q[i] - lim[i].upper) < 1e-9) : [];
    problems.push(M.problem('ik.unreached', 'warning', `Inverse kinematics did not reach the target: ${er.pos.toExponential(2)} m and ${er.rot.toExponential(2)} rad short after ${it} iterations` + (atLimit.length ? `; joint ${atLimit.join(', ')} is at its limit.` : '. The target may be out of reach or at a singularity.'), [{ id: link }], { found: er, req: { pos: tol, rot: rotTol } }));
  }
  return { ok, q, iterations: it, error: er, problems };
}

/** Closed-form IK of a planar two-link arm (links l1, l2) for a tip at (x, y): both elbow solutions [{q1, q2}] ([] when out of reach). */
export function ikPlanar2(l1, l2, x, y) {
  const r2 = x * x + y * y, c2 = (r2 - l1 * l1 - l2 * l2) / (2 * l1 * l2);
  if (c2 > 1 + 1e-12 || c2 < -1 - 1e-12) return [];
  const out = [];
  for (const s of [1, -1]) {
    const q2 = s * Math.acos(Math.max(-1, Math.min(1, c2))), q1 = Math.atan2(y, x) - Math.atan2(l2 * Math.sin(q2), l1 + l2 * Math.cos(q2));
    out.push({ q1, q2 });
  }
  return out;
}

/** 1-D rest-to-rest move of `delta` with peak speed vmax and acceleration amax: triangular when the distance is too short to reach vmax. */
export function trapezoid(delta, vmax, amax) {
  const s = Math.sign(delta) || 1, d = Math.abs(delta);
  if (!(vmax > 0) || !(amax > 0)) return { T: NaN, at: () => ({ x: NaN, v: NaN, a: NaN }) };
  const ta = vmax / amax, da = vmax * vmax / amax;
  let tA, tC, vp;
  if (d >= da) { tA = ta; vp = vmax; tC = (d - da) / vmax; } else { tA = Math.sqrt(d / amax); vp = amax * tA; tC = 0; }
  const T = 2 * tA + tC;
  return { T, peak: vp * s, at(t) {
    t = Math.min(Math.max(t, 0), T);
    let x, v, a;
    if (t < tA) { a = amax; v = amax * t; x = 0.5 * amax * t * t; }
    else if (t < tA + tC) { a = 0; v = vp; x = 0.5 * amax * tA * tA + vp * (t - tA); }
    else { const u = T - t; a = -amax; v = amax * u; x = d - 0.5 * amax * u * u; }
    return { x: s * x, v: s * v, a: s * a };
  } };
}

/** Cubic with zero end velocities over T: q(t) = q0 + Δ(3τ² − 2τ³). Quintic with zero end velocity and acceleration: Δ(10τ³ − 15τ⁴ + 6τ⁵). */
export function cubic(q0, q1, T) { return { T, at(t) { const u = Math.min(Math.max(t / T, 0), 1), d = q1 - q0; return { x: q0 + d * (3 * u * u - 2 * u ** 3), v: d * (6 * u - 6 * u * u) / T, a: d * (6 - 12 * u) / T / T }; } }; }

export function quintic(q0, q1, T) { return { T, at(t) { const u = Math.min(Math.max(t / T, 0), 1), d = q1 - q0; return { x: q0 + d * (10 * u ** 3 - 15 * u ** 4 + 6 * u ** 5), v: d * (30 * u * u - 60 * u ** 3 + 30 * u ** 4) / T, a: d * (60 * u - 180 * u * u + 120 * u ** 3) / T / T }; } }; }

/**
 * jointMove(robot, qa, qb, {vmax, amax}) → {T, at(t) → {q, qd, qdd}, problems}. All joints start and stop together: the slowest joint sets T.
 * vmax / amax may be numbers or arrays; the default velocity comes from each joint's limit.velocity, else 1.
 */
export function jointMove(robot, qa, qb, { vmax, amax = 2 } = {}) {
  const n = robot.dof, vm = (i) => Array.isArray(vmax) ? vmax[i] : (vmax ?? robot.dofJoints[i].limit?.velocity ?? 1), am = (i) => Array.isArray(amax) ? amax[i] : amax;
  const prof = qa.map((a, i) => trapezoid(qb[i] - a, vm(i), am(i))), T = Math.max(0, ...prof.map((p) => p.T));
  const problems = [], lim = robot.limits();
  qb.forEach((v, i) => { if (v < lim[i].lower || v > lim[i].upper) problems.push(M.problem('limit', 'warning', `Target for joint "${robot.names[i]}" (${v}) is outside its limit.`, [{ id: robot.names[i] }])); });
  // every joint follows the slowest joint's profile scaled to its own distance, so all start and stop together
  const si = prof.reduce((b, p, i) => (p.T > prof[b].T ? i : b), 0), dS = qb[si] - qa[si];
  return { T, problems, at(t) {
    const s = T > 0 && dS !== 0 ? prof[si].at(t) : { x: 0, v: 0, a: 0 }, g = dS !== 0 ? 1 / dS : 0;
    return { q: qa.map((a, i) => a + (qb[i] - a) * s.x * g), qd: qa.map((a, i) => (qb[i] - a) * s.v * g), qdd: qa.map((a, i) => (qb[i] - a) * s.a * g) };
  } };
}

/** Straight line in space for `link` from the current pose to point `pB`, n samples through IK (each seeded from the last). → {q[][], ok, problems} */
export function linePath(robot, link, q0, pB, { n = 20, R, local, base } = {}) {
  const T0 = robot.pose(q0, link, base), pA = M.T_apply(T0, local || [0, 0, 0]), q = [q0.slice()], problems = [];
  let cur = q0.slice(), ok = true;
  for (let k = 1; k <= n; k++) {
    const p = M.add(pA, M.scale(M.sub(pB, pA), k / n)), r = ik(robot, link, R ? { p, R } : { p }, cur, { local, base });
    if (!r.ok) { ok = false; problems.push(...r.problems.map((x) => ({ ...x, id: x.id + '@' + k, text: x.text + ` (step ${k} of ${n} along the line)` }))); break; }
    cur = r.q; q.push(cur);
  }
  return { q, ok, problems };
}

/** Peak joint torques of a move against each joint's effort limit. → {peak[], problems} */
export function checkTorques(robot, move, { samples = 50, gravity } = {}) {
  const peak = new Array(robot.dof).fill(0), problems = [];
  for (let k = 0; k <= samples; k++) { const s = move.at(move.T * k / samples), tau = rnea(robot, s.q, s.qd, s.qdd, { gravity }); tau.forEach((v, i) => { peak[i] = Math.max(peak[i], Math.abs(v)); }); }
  robot.dofJoints.forEach((j, i) => { const e = j.limit?.effort; if (e && peak[i] > e) problems.push(M.problem('effort', 'warning', `Joint "${j.name}" needs a peak of ${peak[i].toFixed(2)} (N or N·m) on this move; its limit is ${e}.`, [{ id: j.name }], { found: peak[i], req: e })); });
  return { peak, problems };
}