#include "motion.h"
#include "buttons.h"
#include <math.h>

// Tuning
// Stillness / bias tracking: the gyro is "at rest" when every axis stays inside
// STILL_PP_LSB peak-to-peak for a whole window. Hand tremor alone is far above this,
// so the bias only updates when the remote is put down.
static const int   STILL_WIN          = 104;                  // samples (0.5 s)
static const int   STILL_PP_LSB       = 16;                   // 0.56 dps
static const int   STILL_WINS_UPDATE  = 2;                    // 1 s at rest before updating
static const int   STILL_WINS_JUMP    = 6;                    // 3 s at rest if the bias moved a lot
static const float BIAS_JUMP_RAD      = 0.5f * (PI / 180.0f);
static const float BIAS_UPDATE_K      = 0.3f;

// Gravity estimate: gyro-propagated, pulled toward the accelerometer with this
// time constant, only while |a| is close to 1 g (no hand acceleration).
static const float GRAV_TAU_S         = 1.0f;
static const float GRAV_GATE_G        = 0.15f;
static const float ROLL_MIN_HORIZ     = 0.25f;  // below this the remote points ~straight up/down
static const float ROLL_REANCHOR_UP   = 0.5f;   // button face within 60 deg of up: gravity is unambiguous

// Speed-adaptive smoothing (One Euro style): cutoff = MIN + BETA * speed.
static const float SMOOTH_MIN_HZ      = 1.5f;
static const float SMOOTH_BETA        = 20.0f;  // Hz per rad/s
static const float SMOOTH_MAX_HZ      = 40.0f;

static const float IDLE_RATE_RAD      = 0.06f;  // sleep/idle detection, raw rate
static const unsigned long CLICK_FREEZE_MS = 120;  // hold cursor around button edges

static const float TEST_RATE_RAD      = 0.5f;
static const int   TEST_SEG_SAMPLES   = (int)IMU_ODR_HZ;  // 1 s per segment

// State
static float bias[3]   = {0, 0, 0};
static bool  biasValid = false;
static float up[3]     = {0, 0, 1};
static bool  upSeeded  = false;
static float rollC = 1.0f, rollS = 0.0f;
static bool  rollFromGravity = true;  // next roll is taken from gravity as-is (no continuity)
static float filtX = 0, filtY = 0;  // smoothed yaw/pitch, rad/s
static float accX  = 0, accY  = 0;  // cursor counts not yet sent
static bool  moved = false;

static int     winN = 0;
static int32_t winSum[3];
static int16_t winMin[3], winMax[3];
static int     stillWins = 0;

static int testSample = -1;
static MotionDebug dbg = {};

static void resetWindow() {
  winN = 0;
  for (int i = 0; i < 3; i++) { winSum[i] = 0; winMin[i] = 32767; winMax[i] = -32768; }
}

static void publishBiasStats() {
  #ifdef FEATURE_TELEMETRY
    statBiasRms = sqrtf((bias[0]*bias[0] + bias[1]*bias[1] + bias[2]*bias[2]) / 3.0f);
  #endif
}

void motionReset() {
  filtX = filtY = 0;
  accX = accY = 0;
  upSeeded = false;
  rollFromGravity = true;
  stillWins = 0;
  resetWindow();
}

bool motionApplyCalibration(const float b[3], const float a[3], bool still) {
  bool apply = still || !biasValid;
  if (apply) {
    for (int i = 0; i < 3; i++) bias[i] = b[i];
    biasValid = still;
    publishBiasStats();
  }
  motionReset();
  float an = sqrtf(a[0]*a[0] + a[1]*a[1] + a[2]*a[2]);
  if (an > 0.5f) {
    for (int i = 0; i < 3; i++) up[i] = a[i] / an;
    upSeeded = true;
  }
  return apply;
}

// Continuous bias tracking while the remote is at rest.
static void trackBias(const ImuRaw& r) {
  const int16_t g[3] = {r.gx, r.gy, r.gz};
  for (int i = 0; i < 3; i++) {
    winSum[i] += g[i];
    if (g[i] < winMin[i]) winMin[i] = g[i];
    if (g[i] > winMax[i]) winMax[i] = g[i];
  }
  if (++winN < STILL_WIN) return;

  int pp = 0;
  float mean[3], jump = 0;
  for (int i = 0; i < 3; i++) {
    pp = max(pp, winMax[i] - winMin[i]);
    mean[i] = (float)winSum[i] / winN * GYRO_RAD_PER_LSB;
    jump = max(jump, fabsf(mean[i] - bias[i]));
  }
  resetWindow();

  bool still = pp <= STILL_PP_LSB;
  dbg.still = still;
  stillWins = still ? stillWins + 1 : 0;
  if (stillWins >= STILL_WINS_UPDATE) rollFromGravity = true;  // put down: re-take roll from gravity
  int needed = (biasValid && jump > BIAS_JUMP_RAD) ? STILL_WINS_JUMP : STILL_WINS_UPDATE;
  if (stillWins < needed) return;

  float k = biasValid ? BIAS_UPDATE_K : 1.0f;
  for (int i = 0; i < 3; i++) bias[i] += k * (mean[i] - bias[i]);
  biasValid = true;
  calTempC = cachedTempC;  // temperature compensation is relative to the last bias update
  dbg.biasUpdates++;
  publishBiasStats();
}

static void updateGravity(const float w[3], const float a[3]) {
  float an = sqrtf(a[0]*a[0] + a[1]*a[1] + a[2]*a[2]);
  if (!upSeeded) {
    if (an < 0.5f) return;
    for (int i = 0; i < 3; i++) up[i] = a[i] / an;
    upSeeded = true;
    return;
  }
  // A world-fixed vector seen from the body rotates as du/dt = u x w.
  float ux = up[0] + (up[1]*w[2] - up[2]*w[1]) * IMU_DT;
  float uy = up[1] + (up[2]*w[0] - up[0]*w[2]) * IMU_DT;
  float uz = up[2] + (up[0]*w[1] - up[1]*w[0]) * IMU_DT;

  float err = fabsf(an - 1.0f);
  if (an > 0.1f && err < GRAV_GATE_G) {
    float k = (IMU_DT / GRAV_TAU_S) * (1.0f - err / GRAV_GATE_G);
    ux += k * (a[0]/an - ux);
    uy += k * (a[1]/an - uy);
    uz += k * (a[2]/an - uz);
  }
  float n = sqrtf(ux*ux + uy*uy + uz*uz);
  if (n > 1e-6f) { up[0] = ux/n; up[1] = uy/n; up[2] = uz/n; }
}

void motionProcess(const ImuRaw& r, unsigned long nowMs) {
  dbg.samples++;

  #ifdef FEATURE_AUTO_RECAL
    if (cfg.featureFlags & FLAG_AUTO_RECAL_ENABLED) trackBias(r);
  #endif

  float w[3] = { r.gx * GYRO_RAD_PER_LSB - bias[0],
                 r.gy * GYRO_RAD_PER_LSB - bias[1],
                 r.gz * GYRO_RAD_PER_LSB - bias[2] };
  #ifdef FEATURE_TEMP_COMPENSATION
    if (cfg.featureFlags & FLAG_TEMP_COMP_ENABLED) {
      float corr = TEMP_COMP_COEFF_DPS_C * (cachedTempC - calTempC) * (PI / 180.0f);
      w[1] -= corr; w[2] -= corr;
    }
  #endif
  const float a[3] = { r.ax * ACCEL_G_PER_LSB, r.ay * ACCEL_G_PER_LSB, r.az * ACCEL_G_PER_LSB };
  updateGravity(w, a);

  // Roll compensation. Body frame in the enclosure: +X toward the wrist (tip is -X),
  // +Y right, +Z out of the button face. Twisting the remote is a rotation about X,
  // i.e. of the (gy, gz) plane. Rotate it back using gravity:
  //   yaw   = w . n, n = (0, uy, uz)/|.|   (de-rolled "up")
  //   pitch = w . r, r = (0, uz, -uy)/|.|  (horizontal, across the pointing axis)
  // Held flat this is exactly yaw=gz, pitch=gy. Roll rate (gx) never reaches the cursor.
  //
  // Gravity alone can't tell "tilted back past vertical" from "held upside down":
  // both put the button face below the horizon. Stay on the branch closest to the
  // previous roll, so tilting through vertical never flips the axes, while a real
  // twist (always continuous) is still followed. The branch is re-taken from gravity
  // whenever tracking restarts, the remote is put down, or the button face points
  // roughly up: that pose is never ambiguous, so whatever happened during odd
  // movements, a normal grip always tracks normally again.
  // Near vertical gravity says nothing about twist, so there the twist is followed by
  // integrating the gyro's rate around the pointing axis; otherwise spinning the
  // remote while it points up would leave it on the wrong branch afterwards.
  float horiz = sqrtf(up[1]*up[1] + up[2]*up[2]);
  if (horiz > ROLL_MIN_HORIZ) {
    float c = up[2] / horiz, s = up[1] / horiz;
    bool anchored = rollFromGravity || up[2] > ROLL_REANCHOR_UP;
    if (!anchored && c*rollC + s*rollS < 0.0f) { c = -c; s = -s; }
    rollC = c; rollS = s;
    rollFromGravity = false;
  } else {
    float d = w[0] * IMU_DT, cd = cosf(d), sd = sinf(d);
    float c = rollC*cd - rollS*sd, s = rollS*cd + rollC*sd;
    float n = sqrtf(c*c + s*s);
    rollC = c / n; rollS = s / n;
  }
  float yaw   = w[1]*rollS + w[2]*rollC;
  float pitch = w[1]*rollC - w[2]*rollS;

  if (testSample >= 0) {
    int seg = testSample / TEST_SEG_SAMPLES;
    yaw = pitch = 0;
    if      (seg == 0) yaw   = -TEST_RATE_RAD;  // cursor right
    else if (seg == 1) yaw   =  TEST_RATE_RAD;  // cursor left
    else if (seg == 2) pitch = -TEST_RATE_RAD;  // cursor down
    else if (seg == 3) pitch =  TEST_RATE_RAD;  // cursor up
    if (++testSample >= 5 * TEST_SEG_SAMPLES) testSample = -1;  // 5th second lets the filter settle
  }

  dbg.yaw = yaw; dbg.pitch = pitch;
  if (yaw*yaw + pitch*pitch > IDLE_RATE_RAD*IDLE_RATE_RAD || testSample >= 0) moved = true;

  // Speed-adaptive low-pass: tremor and sensor noise are removed at low speed,
  // while fast movements pass with almost no lag.
  float sPrev = sqrtf(filtX*filtX + filtY*filtY);
  float fc = min(SMOOTH_MIN_HZ + SMOOTH_BETA * sPrev, SMOOTH_MAX_HZ);
  float alpha = 1.0f / (1.0f + 1.0f / (2.0f * PI * fc * IMU_DT));
  filtX += alpha * (yaw   - filtX);
  filtY += alpha * (pitch - filtY);

  // Radial soft dead zone: gain s^2/(s^2+dz^2) is continuous (no jump when motion
  // starts) and acts on the speed, so diagonal motion keeps its direction.
  float s  = sqrtf(filtX*filtX + filtY*filtY);
  float dz = cfg.deadZone;
  float g  = (s > 1e-9f) ? (s*s) / (s*s + dz*dz) : 0.0f;
  float sv = s * g;
  dbg.speed = sv;
  if (sv < 1e-9f) return;

  // Curve and acceleration act on speed, expressed as counts per 10 ms so existing
  // settings keep their feel, then scale both axes equally.
  float dref  = cfg.sensitivity * sv * 0.01f;
  float m     = applyAcceleration(applyCurve(dref));
  float scale = (m / dref) * cfg.sensitivity * IMU_DT * g;
  float cx = -filtX * scale;
  float cy = -filtY * scale;
  if (cfg.axisFlip & 0x01) cx = -cx;
  if (cfg.axisFlip & 0x02) cy = -cy;

  #ifdef FEATURE_PHYSICAL_BUTTONS
    // Pressing a button twists the remote slightly; hold the cursor still around the click.
    if (btnLastEdgeMs && (nowMs - btnLastEdgeMs) < CLICK_FREEZE_MS && testSample < 0) return;
  #endif
  accX += cx; accY += cy;
}

bool motionTakeCounts(int8_t& dx, int8_t& dy) {
  int ix = constrain((int)accX, -127, 127);
  int iy = constrain((int)accY, -127, 127);
  if (ix == 0 && iy == 0) return false;
  accX -= ix; accY -= iy;
  dx = (int8_t)ix; dy = (int8_t)iy;
  return true;
}

void motionReturnCounts(int8_t dx, int8_t dy) { accX += dx; accY += dy; }

bool motionConsumeMoved() { bool m = moved; moved = false; return m; }

void motionStartTest() { testSample = 0; filtX = filtY = 0; accX = accY = 0; }
bool motionTestActive() { return testSample >= 0; }

const MotionDebug& motionDebug() {
  for (int i = 0; i < 3; i++) { dbg.up[i] = up[i]; dbg.bias[i] = bias[i]; }
  dbg.biasValid = biasValid;
  dbg.rollDeg = atan2f(rollS, rollC) * (180.0f / PI);
  return dbg;
}