// 2D similarity transforms and temporal smoothing.

// Least-squares similarity (translation + rotation + uniform scale, 4 DOF)
// mapping P onto Q. Closed form; no iteration.
//
// Deliberately NOT affine or homography: the extra degrees of freedom absorb
// out-of-plane head rotation as shear/perspective and smear it into the mouth.
// Four DOF removes exactly translation, roll and depth-scale, and leaves yaw and
// pitch as a measurable residual.
export function fitSimilarity(P, Q) {
  const n = P.length;
  let pcx = 0, pcy = 0, qcx = 0, qcy = 0;
  for (let i = 0; i < n; i++) {
    pcx += P[i].x; pcy += P[i].y;
    qcx += Q[i].x; qcy += Q[i].y;
  }
  pcx /= n; pcy /= n; qcx /= n; qcy /= n;

  let a = 0, b = 0, norm = 0;
  for (let i = 0; i < n; i++) {
    const px = P[i].x - pcx, py = P[i].y - pcy;
    const qx = Q[i].x - qcx, qy = Q[i].y - qcy;
    a += px * qx + py * qy;   // dot
    b += px * qy - py * qx;   // cross
    norm += px * px + py * py;
  }
  const theta = Math.atan2(b, a);
  const s = norm > 1e-12 ? Math.hypot(a, b) / norm : 1;
  const c = Math.cos(theta), sn = Math.sin(theta);
  return {
    s, theta,
    tx: qcx - s * (c * pcx - sn * pcy),
    ty: qcy - s * (sn * pcx + c * pcy),
  };
}

export function applySim(tf, p) {
  const c = Math.cos(tf.theta), sn = Math.sin(tf.theta);
  return {
    x: tf.s * (c * p.x - sn * p.y) + tf.tx,
    y: tf.s * (sn * p.x + c * p.y) + tf.ty,
  };
}

export function applySimAll(tf, pts) {
  return pts.map((p) => applySim(tf, p));
}

// Residual RMS after the fit, in the units of Q. Rises with out-of-plane
// rotation, so it is the signal for "this section is not stabilisable".
export function fitResidual(tf, P, Q) {
  let acc = 0;
  for (let i = 0; i < P.length; i++) {
    const m = applySim(tf, P[i]);
    acc += (m.x - Q[i].x) ** 2 + (m.y - Q[i].y) ** 2;
  }
  return Math.sqrt(acc / P.length);
}

// Generalised Procrustes: the reference is the MEAN rigid configuration over the
// shot, not frame zero, so no single frame's idiosyncrasies get baked into every
// other frame. Three passes is plenty.
export function procrustesMean(framesRigid, iters = 3) {
  let ref = framesRigid[0].map((p) => ({ x: p.x, y: p.y }));
  for (let it = 0; it < iters; it++) {
    const acc = ref.map(() => ({ x: 0, y: 0 }));
    for (const rig of framesRigid) {
      const tf = fitSimilarity(rig, ref);
      const moved = applySimAll(tf, rig);
      for (let i = 0; i < acc.length; i++) { acc[i].x += moved[i].x; acc[i].y += moved[i].y; }
    }
    ref = acc.map((p) => ({ x: p.x / framesRigid.length, y: p.y / framesRigid.length }));
  }
  return ref;
}

function movingAverage(vals, win) {
  if (win <= 1) return vals.slice();
  const half = Math.floor(win / 2), out = new Array(vals.length);
  for (let i = 0; i < vals.length; i++) {
    let acc = 0, cnt = 0;
    for (let j = i - half; j <= i + half; j++) {
      const k = Math.min(vals.length - 1, Math.max(0, j));
      acc += vals[k]; cnt++;
    }
    out[i] = acc / cnt;
  }
  return out;
}

export { movingAverage };

// Smooth the four transform parameters, NEVER the contour. Landmark jitter of a
// pixel is smeared into the mouth by the inverse transform, so the transform is
// where the low-pass belongs; smoothing the contour would destroy the
// performance, which is the entire asset.
// Angles are smoothed as (cos, sin) so wrapping cannot produce a spike.
export function smoothTransforms(tfs, win) {
  const c = movingAverage(tfs.map((t) => Math.cos(t.theta)), win);
  const sn = movingAverage(tfs.map((t) => Math.sin(t.theta)), win);
  const s = movingAverage(tfs.map((t) => t.s), win);
  const tx = movingAverage(tfs.map((t) => t.tx), win);
  const ty = movingAverage(tfs.map((t) => t.ty), win);
  return tfs.map((_, i) => ({
    theta: Math.atan2(sn[i], c[i]), s: s[i], tx: tx[i], ty: ty[i],
  }));
}