#include "track/neck_cal.h"

#include <cmath>
#include <cstdio>

namespace sauna {
namespace {

constexpr double kG = 9.80665;

// Gates. Ballparks from the design discussion (doc §4): a human neck-to-eye
// lever is ~7-12 cm forward, ~10-18 cm up; lateral is anatomy-zero. The
// residual gate is the torso-sway catch: translation of the pivot is
// unmodeled acceleration and inflates the fit error.
constexpr double kMinPhaseRms = 0.8;     // rad/s about the phase axis
constexpr double kMaxLateralMm = 35.0;
constexpr double kMaxResidualRms = 0.45; // m/s^2
constexpr double kMinSamples = 2000;     // ~2 s of motion at 1 kHz

void cross(const double a[3], const double b[3], double out[3]) {
  out[0] = a[1] * b[2] - a[2] * b[1];
  out[1] = a[2] * b[0] - a[0] * b[2];
  out[2] = a[0] * b[1] - a[1] * b[0];
}

// Rotate v by q^-1 (world -> head when q is world<-head).
void rotateInv(const Quat& q, const double v[3], double out[3]) {
  const Quat c{q.w, -q.x, -q.y, -q.z};
  const double qv[3] = {c.x, c.y, c.z};
  double t[3], u[3];
  cross(qv, v, t);
  t[0] *= 2; t[1] *= 2; t[2] *= 2;
  cross(qv, t, u);
  out[0] = v[0] + c.w * t[0] + u[0];
  out[1] = v[1] + c.w * t[1] + u[1];
  out[2] = v[2] + c.w * t[2] + u[2];
}

// Solve N x = v for symmetric 6x6 N, Gaussian elimination, partial pivot.
bool solve6(double N[6][6], double v[6], double x[6]) {
  int idx[6] = {0, 1, 2, 3, 4, 5};
  for (int c = 0; c < 6; c++) {
    int p = c;
    for (int r = c + 1; r < 6; r++)
      if (std::fabs(N[idx[r]][c]) > std::fabs(N[idx[p]][c])) p = r;
    std::swap(idx[c], idx[p]);
    const double piv = N[idx[c]][c];
    if (std::fabs(piv) < 1e-9) return false;  // unobservable
    for (int r = c + 1; r < 6; r++) {
      const double m = N[idx[r]][c] / piv;
      for (int k = c; k < 6; k++) N[idx[r]][k] -= m * N[idx[c]][k];
      v[idx[r]] -= m * v[idx[c]];
    }
  }
  for (int c = 5; c >= 0; c--) {
    double s = v[idx[c]];
    for (int k = c + 1; k < 6; k++) s -= N[idx[c]][k] * x[k];
    x[c] = s / N[idx[c]][c];
  }
  return true;
}

}  // namespace

void NeckCalibrator::feed(const double omegaHead[3], const double accHeadG[3],
                          const Quat& q, double dt, int phase) {
  if (phase == kIdle) return;
  if (dt <= 0.0 || dt > 0.1) {
    // Stream gap — restart the local time axis, derivative skips it.
    tAcc_ += 1.0;
    return;
  }
  tAcc_ += dt;
  Sample s;
  s.t = tAcc_;
  for (int a = 0; a < 3; a++) {
    s.w[a] = omegaHead[a];
    s.f[a] = accHeadG[a] * kG;
  }
  const double upWorld[3] = {0, 1, 0};
  rotateInv(q, upWorld, s.up);
  s.phase = phase;
  samples_.push_back(s);
}

void NeckCalibrator::clear() {
  samples_.clear();
  tAcc_ = 0.0;
}

NeckCalibrator::Result NeckCalibrator::solve(
    const double imuPosHead[3]) const {
  Result res;
  res.samples = samples_.size();
  if (samples_.size() < kMinSamples) {
    res.message = "too few samples — was the headset streaming?";
    return res;
  }

  // Excitation per phase: yaw shakes rotate about head +y, pitch nods
  // about head +x.
  double yawSq = 0, pitchSq = 0;
  uint64_t yawN = 0, pitchN = 0;
  for (const auto& s : samples_) {
    if (s.phase == kYaw) { yawSq += s.w[1] * s.w[1]; yawN++; }
    if (s.phase == kPitch) { pitchSq += s.w[0] * s.w[0]; pitchN++; }
  }
  res.yawRms = yawN ? std::sqrt(yawSq / yawN) : 0.0;
  res.pitchRms = pitchN ? std::sqrt(pitchSq / pitchN) : 0.0;
  if (res.yawRms < kMinPhaseRms || res.pitchRms < kMinPhaseRms) {
    char m[160];
    snprintf(m, sizeof(m),
             "not enough rotation (yaw %.2f / pitch %.2f rad/s RMS, need "
             ">= %.1f) — shake/nod with more energy",
             res.yawRms, res.pitchRms, kMinPhaseRms);
    res.message = m;
    return res;
  }

  // Normal equations over rows  [ [wdot x] + [w x][w x] | I3 ] [r; b] = y,
  // y = f - u(q)*g. Angular acceleration by central difference over +-2
  // samples (1 kHz stream: ~4 ms base), gap-safe via the time axis.
  double N[6][6] = {}, v[6] = {};
  uint64_t rows = 0;
  const auto& S = samples_;
  for (size_t i = 2; i + 2 < S.size(); i++) {
    const double dt4 = S[i + 2].t - S[i - 2].t;
    if (dt4 <= 0.0 || dt4 > 0.05) continue;  // spans a gap — skip
    double wdot[3];
    for (int a = 0; a < 3; a++) wdot[a] = (S[i + 2].w[a] - S[i - 2].w[a]) / dt4;
    const double* w = S[i].w;
    // M = [wdot x] + [w x][w x];  [w x][w x] = w w^T - |w|^2 I.
    const double w2 = w[0] * w[0] + w[1] * w[1] + w[2] * w[2];
    double M[3][3] = {
        {w[0] * w[0] - w2, w[0] * w[1] - wdot[2], w[0] * w[2] + wdot[1]},
        {w[1] * w[0] + wdot[2], w[1] * w[1] - w2, w[1] * w[2] - wdot[0]},
        {w[2] * w[0] - wdot[1], w[2] * w[1] + wdot[0], w[2] * w[2] - w2}};
    double y[3];
    for (int a = 0; a < 3; a++) y[a] = S[i].f[a] - S[i].up[a] * kG;
    // Accumulate A^T A and A^T y with A = [M | I].
    for (int r = 0; r < 3; r++) {
      for (int c = 0; c < 3; c++) {
        double mm = 0;
        for (int k = 0; k < 3; k++) mm += M[k][r] * M[k][c];
        N[r][c] += mm;                       // M^T M
      }
      for (int c = 0; c < 3; c++) {
        N[r][3 + c] += M[c][r];              // M^T I
        N[3 + c][r] += M[c][r];              // I M
      }
      N[3 + r][3 + r] += 1.0;                // I I
      double my = 0;
      for (int k = 0; k < 3; k++) my += M[k][r] * y[k];
      v[r] += my;
      v[3 + r] += y[r];
    }
    rows++;
  }
  if (rows < kMinSamples / 2) {
    res.message = "too few usable samples after gap filtering";
    return res;
  }

  double x[6];
  if (!solve6(N, v, x)) {
    res.message = "solve is degenerate — motion did not observe the lever "
                  "arm (rotate, don't translate)";
    return res;
  }
  for (int a = 0; a < 3; a++) {
    res.rImuM[a] = x[a];
    res.accBias[a] = x[3 + a];
    res.rEyeM[a] = x[a] - imuPosHead[a];
  }
  res.neckForwardMm = -res.rEyeM[2] * 1000.0;
  res.neckUpMm = res.rEyeM[1] * 1000.0;

  // Residual over the same rows.
  double sq = 0;
  uint64_t n = 0;
  for (size_t i = 2; i + 2 < S.size(); i++) {
    const double dt4 = S[i + 2].t - S[i - 2].t;
    if (dt4 <= 0.0 || dt4 > 0.05) continue;
    double wdot[3];
    for (int a = 0; a < 3; a++) wdot[a] = (S[i + 2].w[a] - S[i - 2].w[a]) / dt4;
    const double* w = S[i].w;
    const double r0 = res.rImuM[0], r1 = res.rImuM[1], r2 = res.rImuM[2];
    double wxr[3] = {w[1] * r2 - w[2] * r1, w[2] * r0 - w[0] * r2,
                     w[0] * r1 - w[1] * r0};
    double wxwxr[3], wdotxr[3];
    cross(w, wxr, wxwxr);
    const double rv[3] = {r0, r1, r2};
    cross(wdot, rv, wdotxr);
    for (int a = 0; a < 3; a++) {
      const double pred =
          wdotxr[a] + wxwxr[a] + S[i].up[a] * kG + res.accBias[a];
      const double e = S[i].f[a] - pred;
      sq += e * e;
    }
    n++;
  }
  res.residualRms = n ? std::sqrt(sq / (3.0 * n)) : 0.0;

  // Gates.
  char m[256];
  if (std::fabs(res.rEyeM[0]) * 1000.0 > kMaxLateralMm) {
    snprintf(m, sizeof(m),
             "lateral component %.0f mm (limit %.0f) — torso moved or the "
             "capture was off; redo with the torso still",
             res.rEyeM[0] * 1000.0, kMaxLateralMm);
    res.message = m;
    return res;
  }
  if (res.residualRms > kMaxResidualRms) {
    snprintf(m, sizeof(m),
             "residual %.2f m/s^2 (limit %.2f) — torso sway or translation "
             "during capture; redo seated, shoulders against the backrest",
             res.residualRms, kMaxResidualRms);
    res.message = m;
    return res;
  }
  snprintf(m, sizeof(m),
           "fwd %.0f mm, up %.0f mm, lateral %.0f mm; residual %.2f m/s^2, "
           "yaw %.1f / pitch %.1f rad/s RMS, %llu rows",
           res.neckForwardMm, res.neckUpMm, res.rEyeM[0] * 1000.0,
           res.residualRms, res.yawRms, res.pitchRms,
           (unsigned long long)rows);
  res.message = m;
  res.ok = true;
  return res;
}

}  // namespace sauna