#include "track/ahrs.h"

#include <chrono>
#include <cmath>

namespace sauna {
namespace {

constexpr double kTimecodeHz = 48e6;

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];
}

double norm3(const double v[3]) {
  return std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
}

Quat mul(const Quat& a, const Quat& b) {
  return {a.w * b.w - a.x * b.x - a.y * b.y - a.z * b.z,
          a.w * b.x + a.x * b.w + a.y * b.z - a.z * b.y,
          a.w * b.y - a.x * b.z + a.y * b.w + a.z * b.x,
          a.w * b.z + a.x * b.y - a.y * b.x + a.z * b.w};
}

void normalize(Quat* q) {
  const double n =
      std::sqrt(q->w * q->w + q->x * q->x + q->y * q->y + q->z * q->z);
  if (n <= 0) {
    *q = {};
    return;
  }
  q->w /= n;
  q->x /= n;
  q->y /= n;
  q->z /= n;
}

// Rotate v by q (v_out = q * v * q^-1).
void rotate(const Quat& q, const double v[3], double out[3]) {
  // t = 2 q_vec x v;  out = v + q.w t + q_vec x t
  const double qv[3] = {q.x, q.y, q.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] + q.w * t[0] + u[0];
  out[1] = v[1] + q.w * t[1] + u[1];
  out[2] = v[2] + q.w * t[2] + u[2];
}

Quat conj(const Quat& q) { return {q.w, -q.x, -q.y, -q.z}; }

}  // namespace

void Ahrs::configure(const HmdConfig& cfg, const Params& params) {
  std::lock_guard<std::mutex> lk(mu_);
  p_ = params;
  // R_head_imu columns = IMU basis vectors expressed in head frame.
  // plus_y = plus_z x plus_x (right-handed).
  const FrameCalib& f = cfg.imu.frame;
  double py[3];
  cross(f.plus_z, f.plus_x, py);
  for (int r = 0; r < 3; r++) {
    rHeadImu_[r][0] = f.plus_x[r];
    rHeadImu_[r][1] = py[r];
    rHeadImu_[r][2] = f.plus_z[r];
  }
  for (int a = 0; a < 3; a++) {
    accBias_[a] = cfg.imu.acc_bias[a];
    accScaleCfg_[a] = cfg.imu.acc_scale[a];
    gyroBiasCfg_[a] = cfg.imu.gyro_bias[a];
  }
  resetLocked(/*forceCapture=*/false);
}

void Ahrs::resetLocked(bool forceCapture) {
  q_ = {};
  haveTc_ = false;
  biasSum_[0] = biasSum_[1] = biasSum_[2] = 0;
  biasSqSum_[0] = biasSqSum_[1] = biasSqSum_[2] = 0;
  biasN_ = 0;
  biasTimeAcc_ = 0;
  // Window continuity is broken — escalation streak and any 'b'
  // attestation don't survive a reset (requestBiasCapture re-arms relax
  // after calling here).
  saneRejStreak_ = 0;
  relaxSane_ = false;
  st_.initialized = false;
  if (haveBias_ && !forceCapture) {
    // Bias already known (seed or earlier capture): skip the still window,
    // re-level from the next gravity sample and keep tracking. Turns the
    // hard gap reset from a freeze-for-recapture into an instant resume.
    biasCapturing_ = false;
    needLevel_ = true;
  } else {
    biasCapturing_ = true;
    needLevel_ = false;
  }
}

void Ahrs::requestBiasCapture() {
  std::lock_guard<std::mutex> lk(mu_);
  resetLocked(/*forceCapture=*/true);
  // Explicit request attests stillness: waive the sane gate for this
  // capture (sd gate still enforced). See Params escalation comment.
  relaxSane_ = true;
}

void Ahrs::seedBias(const double biasLsb[3]) {
  std::lock_guard<std::mutex> lk(mu_);
  for (int a = 0; a < 3; a++) {
    gyroBias_[a] = biasLsb[a];
    st_.gyroBiasLsb[a] = biasLsb[a];
  }
  haveBias_ = true;
  if (biasCapturing_) {
    biasCapturing_ = false;
    needLevel_ = true;
    biasSum_[0] = biasSum_[1] = biasSum_[2] = 0;
    biasSqSum_[0] = biasSqSum_[1] = biasSqSum_[2] = 0;
    biasN_ = 0;
    biasTimeAcc_ = 0;
  }
}

bool Ahrs::takeBiasRefresh(double outLsb[3]) {
  std::lock_guard<std::mutex> lk(mu_);
  if (!biasDirty_) return false;
  for (int a = 0; a < 3; a++) outLsb[a] = gyroBias_[a];
  biasDirty_ = false;
  return true;
}

bool Ahrs::takeBiasCaptureEvent(BiasCaptureEvent* out) {
  std::lock_guard<std::mutex> lk(mu_);
  if (!haveCaptureEvent_) return false;
  *out = captureEvent_;
  haveCaptureEvent_ = false;
  return true;
}

bool Ahrs::saneOkLocked(double meanLsb, int axis) const {
  if (std::fabs(meanLsb - gyroBiasCfg_[axis]) <= p_.biasSaneLsb) return true;
  if (haveBias_ && std::fabs(meanLsb - gyroBias_[axis]) <= p_.biasSaneLsb)
    return true;
  return false;
}

void Ahrs::levelFromGravityLocked() {
  // Initialize orientation from gravity: find q aligning measured accel
  // dir (head frame) with world +Y, yaw = 0 by construction (shortest-arc
  // rotation).
  double up[3] = {0, 1, 0};
  double v[3] = {st_.gravityHead[0], st_.gravityHead[1], st_.gravityHead[2]};
  double axis[3];
  cross(v, up, axis);
  const double c = v[0] * up[0] + v[1] * up[1] + v[2] * up[2];
  Quat q;
  if (c < -0.9999) {
    q = {0, 0, 0, 1};  // 180° about Z (any horizontal axis works)
  } else {
    q.w = 1.0 + c;
    q.x = axis[0];
    q.y = axis[1];
    q.z = axis[2];
    normalize(&q);
  }
  q_ = q;
  needLevel_ = false;
  st_.initialized = true;
}

void Ahrs::update(const ImuSample& s) {
  std::lock_guard<std::mutex> lk(mu_);
  st_.samples++;
  lastSampleHostNs_ = std::chrono::duration_cast<std::chrono::nanoseconds>(
                          std::chrono::steady_clock::now().time_since_epoch())
                          .count();

  double dt = 0.0;
  if (haveTc_) {
    dt = (uint32_t)(s.timecode - lastTc_) / kTimecodeHz;
    if (dt > p_.gapResetSec) {
      if (st_.initialized && dt < p_.hardResetSec) {
        // Micro-gap (sub-second stream stall, not a reconnect): keep the
        // orientation and the captured bias, just don't integrate across
        // the hole. The old full reset here zeroed the pose to identity
        // and froze tracking for a fresh still-window bias capture — a
        // brief USB hiccup became a multi-second frozen-then-snap episode
        // (M3 field report: "rubber-bandy"). Tilt error accrued during
        // the gap pulls back in through the accel term; yaw error is
        // accepted (recenter owns yaw).
        st_.softRecoveries++;
        dt = 0.0;  // resume integration from the next sample
      } else {
        // Long gap (reconnect/silence) — re-level rather than integrate junk.
        // With a known bias this resumes from the next gravity sample
        // (resetLocked keeps it); without one it falls back to the still
        // window.
        st_.resets++;
        resetLocked(/*forceCapture=*/false);
      }
    }
  }
  lastTc_ = s.timecode;
  haveTc_ = true;

  // Accel in g (config bias/scale applied; identity/zero on known units, so
  // the exact libsurvive semantics don't bite — assumed scale-then-bias).
  double accImu[3], accHead[3];
  for (int a = 0; a < 3; a++)
    accImu[a] = (s.accel[a] / p_.accelScale) * accScaleCfg_[a] - accBias_[a];
  for (int r = 0; r < 3; r++)
    accHead[r] = rHeadImu_[r][0] * accImu[0] + rHeadImu_[r][1] * accImu[1] +
                 rHeadImu_[r][2] * accImu[2];
  const double accMag = norm3(accHead);
  st_.accelMagG = accMag;
  if (accMag > 1e-6) {
    st_.gravityHead[0] = accHead[0] / accMag;
    st_.gravityHead[1] = accHead[1] / accMag;
    st_.gravityHead[2] = accHead[2] / accMag;
  }

  if (biasCapturing_) {
    for (int a = 0; a < 3; a++) {
      biasSum_[a] += s.gyro[a];
      biasSqSum_[a] += (double)s.gyro[a] * s.gyro[a];
    }
    biasN_++;
    biasTimeAcc_ += dt;
    if (biasTimeAcc_ >= p_.biasWindowSec && biasN_ >= 100) {
      double mean[3], sd[3];
      bool sdFail = false, saneFail = false;
      for (int a = 0; a < 3; a++) {
        mean[a] = biasSum_[a] / biasN_;
        sd[a] = std::sqrt(
            std::fmax(0.0, biasSqSum_[a] / biasN_ - mean[a] * mean[a]));
        if (sd[a] > p_.biasStillLsb) sdFail = true;
        // Mean sanity too: a slow smooth rotation (worn-at-launch head
        // drift) passes the stddev gate but pollutes the mean.
        if (!saneOkLocked(mean[a], a)) saneFail = true;
      }
      bool accept = !sdFail && (!saneFail || relaxSane_);
      BiasCaptureEvent::Kind kind = (relaxSane_ && saneFail)
                                        ? BiasCaptureEvent::Kind::kAcceptedRelaxed
                                        : BiasCaptureEvent::Kind::kRejected;
      if (!accept && !sdFail && saneFail) {
        // Sane-only reject: count consecutive windows whose means agree —
        // a genuine large rest bias holds steady, head drift does not.
        bool consistent = saneRejStreak_ > 0;
        for (int a = 0; a < 3 && consistent; a++)
          if (std::fabs(mean[a] - saneRejMean_[a]) >
              p_.biasEscalateConsistLsb)
            consistent = false;
        saneRejStreak_ = consistent ? saneRejStreak_ + 1 : 1;
        for (int a = 0; a < 3; a++) saneRejMean_[a] = mean[a];
        if (saneRejStreak_ >= p_.biasEscalateWindows) {
          accept = true;
          kind = BiasCaptureEvent::Kind::kAcceptedEscalated;
        }
      } else if (sdFail) {
        saneRejStreak_ = 0;  // motion breaks the consistency argument
      }
      if (accept) {
        for (int a = 0; a < 3; a++) {
          gyroBias_[a] = mean[a];
          st_.gyroBiasLsb[a] = mean[a];
        }
        biasCapturing_ = false;
        haveBias_ = true;
        biasDirty_ = true;  // freshly measured — worth persisting
        levelFromGravityLocked();
        if (kind != BiasCaptureEvent::Kind::kRejected) {
          // Escape-path accepts are worth telling the user about; the
          // in-gates accept stays quiet as before.
          captureEvent_ = {kind, {mean[0], mean[1], mean[2]},
                           {sd[0], sd[1], sd[2]}, false, true,
                           saneRejStreak_};
          haveCaptureEvent_ = true;
        }
        relaxSane_ = false;
        saneRejStreak_ = 0;
      } else {
        captureEvent_ = {BiasCaptureEvent::Kind::kRejected,
                         {mean[0], mean[1], mean[2]},
                         {sd[0], sd[1], sd[2]}, sdFail, saneFail,
                         saneRejStreak_};
        haveCaptureEvent_ = true;
        // Not accepted — slide the window.
        for (int a = 0; a < 3; a++) {
          biasSum_[a] = 0;
          biasSqSum_[a] = 0;
        }
        biasN_ = 0;
        biasTimeAcc_ = 0;
      }
    }
    return;
  }

  if (needLevel_) {
    // Bias known (seed or pre-reset capture), orientation pending: level
    // from the first sane gravity sample and track immediately — no still
    // window (instant start: users don the headset before launching).
    // Strong linear acceleration just defers leveling to the next sample;
    // residual tilt pulls in through the accel term.
    if (accMag > 0.5 && accMag < 1.5) levelFromGravityLocked();
    return;
  }

  if (dt <= 0.0) return;

  // Opportunistic background bias refresh (M4 step 1): keep accumulating
  // still-window stats while tracking; an accepted window updates the bias
  // in place (delta is ~LSBs — invisible mid-run) and flags it for
  // persistence. Low stddev alone also passes a slow CONSTANT rotation, so
  // the mean must additionally be plausible as a rest bias. Sane ref is
  // either the config bias OR the trusted bias (saneOkLocked): config-only
  // dead-zoned refreshes entirely on a unit whose genuine rest bias
  // exceeds the gate; the either-ref keeps the self-heal (a true still
  // window near config always accepts and pulls a polluted value back).
  for (int a = 0; a < 3; a++) {
    biasSum_[a] += s.gyro[a];
    biasSqSum_[a] += (double)s.gyro[a] * s.gyro[a];
  }
  biasN_++;
  biasTimeAcc_ += dt;
  if (biasTimeAcc_ >= p_.biasWindowSec && biasN_ >= 100) {
    bool accept = true;
    double mean[3];
    for (int a = 0; a < 3; a++) {
      mean[a] = biasSum_[a] / biasN_;
      const double sd = std::sqrt(
          std::fmax(0.0, biasSqSum_[a] / biasN_ - mean[a] * mean[a]));
      if (sd > p_.biasStillLsb || !saneOkLocked(mean[a], a)) accept = false;
    }
    if (accept) {
      double maxDelta = 0.0;
      for (int a = 0; a < 3; a++) {
        maxDelta = std::fmax(maxDelta, std::fabs(mean[a] - gyroBias_[a]));
        gyroBias_[a] = mean[a];
        st_.gyroBiasLsb[a] = mean[a];
      }
      st_.biasRefreshes++;
      if (maxDelta > p_.biasDirtyLsb) biasDirty_ = true;
    }
    biasSum_[0] = biasSum_[1] = biasSum_[2] = 0;
    biasSqSum_[0] = biasSqSum_[1] = biasSqSum_[2] = 0;
    biasN_ = 0;
    biasTimeAcc_ = 0;
  }

  // Gyro rad/s in head frame (captured still-bias + config bias, then scale).
  double gyroImu[3], omega[3];
  for (int a = 0; a < 3; a++)
    gyroImu[a] = (s.gyro[a] - gyroBias_[a] - gyroBiasCfg_[a]) * p_.gyroScale;
  for (int r = 0; r < 3; r++)
    omega[r] = rHeadImu_[r][0] * gyroImu[0] + rHeadImu_[r][1] * gyroImu[1] +
               rHeadImu_[r][2] * gyroImu[2];
  st_.omegaRadS = norm3(omega);  // motion magnitude before any correction

  // Neck-cal tap: raw physical rate + accel + orientation, pre-correction.
  if (tap_) tap_(omega, accHead, q_, dt);

  // Mahony proportional term: error = measured gravity dir x estimated
  // gravity dir (both in head frame). Skip when |a| is far from 1 g
  // (linear acceleration corrupts the reference). Additionally fade the
  // gain out with rotation rate: during head motion the accelerometer
  // carries centripetal/tangential components inside the magnitude gate,
  // and the resulting correction tug reads as hitching once a neck model
  // amplifies orientation into translation. At rest full kp (drift gate
  // owns rest behavior); above corrFadeFullRadS no correction. Gates are
  // params (M4 step 3 tightening: 0.85-1.15 g, fade out by 0.5 rad/s —
  // the old 0.5-1.5 g / 1.0 rad/s admitted 0.13-0.18 rad/s tugs during
  // ordinary motion). Tilt error accrued while gated pulls in at rest.
  if (accMag > p_.corrGateMinG && accMag < p_.corrGateMaxG) {
    const double rate = st_.omegaRadS;
    const double fade = std::fmax(0.0, 1.0 - rate / p_.corrFadeFullRadS);
    if (fade > 0.0) {
      double upWorld[3] = {0, 1, 0};
      double vEst[3];
      rotate(conj(q_), upWorld, vEst);  // world up in head frame per estimate
      double e[3];
      const double vMeas[3] = {st_.gravityHead[0], st_.gravityHead[1],
                               st_.gravityHead[2]};
      cross(vMeas, vEst, e);
      double gain = p_.kp * fade;
      double corr = gain * norm3(e);
      if (corr > p_.corrMaxRadS) {  // absolute tug ceiling (see Params)
        gain *= p_.corrMaxRadS / corr;
        corr = p_.corrMaxRadS;
      }
      for (int a = 0; a < 3; a++) omega[a] += gain * e[a];
      if (corr > maxCorr_) maxCorr_ = corr;
    }
  }

  lastOmega_[0] = omega[0];
  lastOmega_[1] = omega[1];
  lastOmega_[2] = omega[2];

  // q_dot = 0.5 * q * (0, omega_head); integrate, renormalize.
  const Quat w{0, omega[0], omega[1], omega[2]};
  Quat dq = mul(q_, w);
  q_.w += 0.5 * dq.w * dt;
  q_.x += 0.5 * dq.x * dt;
  q_.y += 0.5 * dq.y * dt;
  q_.z += 0.5 * dq.z * dt;
  normalize(&q_);
}

Quat Ahrs::pose(double predictSec) const {
  std::lock_guard<std::mutex> lk(mu_);
  if (predictSec <= 0.0) return q_;
  if (predictSec > 0.05) predictSec = 0.05;
  // q * exp(0.5 * w_body * dt)
  const double wx = lastOmega_[0], wy = lastOmega_[1], wz = lastOmega_[2];
  const double mag = std::sqrt(wx * wx + wy * wy + wz * wz);
  if (mag < 1e-9) return q_;
  const double half = 0.5 * mag * predictSec;
  const double s = std::sin(half) / mag;
  Quat d{std::cos(half), wx * s, wy * s, wz * s};
  Quat out = mul(q_, d);
  normalize(&out);
  return out;
}

Quat Ahrs::poseAtHostNs(int64_t hostTimeNs, double* outAgeSec) const {
  std::lock_guard<std::mutex> lk(mu_);
  double predictSec = 0.0;
  if (lastSampleHostNs_ != 0)
    predictSec = (double)(hostTimeNs - lastSampleHostNs_) * 1e-9;
  if (predictSec < 0.0) predictSec = 0.0;
  if (predictSec > 0.05) predictSec = 0.05;
  if (outAgeSec) *outAgeSec = predictSec;
  if (predictSec <= 0.0) return q_;
  const double wx = lastOmega_[0], wy = lastOmega_[1], wz = lastOmega_[2];
  const double mag = std::sqrt(wx * wx + wy * wy + wz * wz);
  if (mag < 1e-9) return q_;
  const double half = 0.5 * mag * predictSec;
  const double s = std::sin(half) / mag;
  Quat d{std::cos(half), wx * s, wy * s, wz * s};
  Quat out = mul(q_, d);
  normalize(&out);
  return out;
}

void Ahrs::recenter() {
  std::lock_guard<std::mutex> lk(mu_);
  // Current head -Z (forward) in world, projected to horizontal -> yaw;
  // premultiply by the inverse yaw about world +Y.
  double fwdHead[3] = {0, 0, -1};
  double fwdWorld[3];
  rotate(q_, fwdHead, fwdWorld);
  const double yaw = std::atan2(-fwdWorld[0], -fwdWorld[2]);
  const Quat qInvYaw{std::cos(-yaw / 2), 0, std::sin(-yaw / 2), 0};
  q_ = mul(qInvYaw, q_);
  normalize(&q_);
}

Ahrs::Status Ahrs::status() const {
  std::lock_guard<std::mutex> lk(mu_);
  return st_;
}

void Ahrs::setSampleTap(SampleTap tap) {
  std::lock_guard<std::mutex> lk(mu_);
  tap_ = std::move(tap);
}

double Ahrs::takeMaxCorrectionRadS() {
  std::lock_guard<std::mutex> lk(mu_);
  const double out = maxCorr_;
  maxCorr_ = 0.0;
  return out;
}

}  // namespace sauna