// AHRS — own Mahony-style orientation filter (ADR locked decision 4).
//
// Gyro integration + accelerometer gravity correction (proportional term
// only, conservative gain, no magnetometer — yaw drift is accepted, the
// recenter UX owns it). Runs in the HEAD frame: raw IMU samples are mapped
// through the config imu block (plus_x/plus_z basis) before integration.
//
// Facts this implementation encodes:
//  - Accel 4096 LSB/g (S1). Gyro scale from gate-2 measurement (config
//    carries NO gyro_scale on this unit; caller supplies the value).
//  - timecode = 48 MHz uint32, wraps ~89.5 s — dt via uint32 subtraction.
//  - Config gyro_bias is zero but the real rest bias is ~1-4 LSB (~13°/min
//    yaw at 0.001 rad/s/LSB), so the filter captures bias during an initial
//    still window; recapturable on demand (stillness validated).
//  - Instant start (M4 step 1): a persisted bias can be seeded via
//    seedBias() — the filter then levels from the first gravity sample and
//    tracks immediately, no still-window wait (users don the headset before
//    launching). While running, every accepted still window refreshes the
//    bias opportunistically; takeBiasRefresh() hands the app a value to
//    persist. With a known bias, hard gap resets also re-level instantly
//    instead of freezing for a recapture.
//  - M1 resilience carry-over: sample gap > ~0.5 s resets the filter
//    (re-levels from accel; device reconnect must not wedge tracking).
//
// Conventions: world = +Y up, -Z forward (OpenVR). pose() returns q such
// that v_world = q * v_head * q^-1. Recenter zeroes yaw only.

#pragma once

#include <cstdint>
#include <functional>
#include <mutex>

#include "calib/hmd_config.h"
#include "device/tundra_imu.h"

namespace sauna {

struct Quat {
  double w = 1.0, x = 0.0, y = 0.0, z = 0.0;
};

class Ahrs {
 public:
  struct Params {
    // Gate-2 measurement (gyro_cal_vr vs SteamVR, LHR-1F8E25F1): gyro full
    // scale is the standard MEMS ±2000 °/s over int16; measured +0.07% off
    // nominal with 0.44% per-axis spread.
    double gyroScale = 0.0010652644360316951;  // 2000*(pi/180)/32768 rad/s/LSB
    double accelScale = 4096.0;   // LSB/g (S1)
    double kp = 0.5;              // Mahony proportional gain
    double gapResetSec = 0.5;     // sample gap that interrupts integration
    double hardResetSec = 5.0;    // gap that forces the full re-level path
    double biasWindowSec = 1.0;   // startup still-window for gyro bias
    double biasStillLsb = 3.0;    // max gyro stddev to accept the window
    // Window acceptance sanity (startup capture AND background refresh):
    // reject a window whose mean is implausibly far from the CONFIG rest
    // bias (~zero; real rest bias is 1-4 LSB) — low stddev alone passes a
    // slow constant rotation, which pollutes the mean (field finding: a
    // worn-at-launch capture stored a drifted bias => steady-state tilt as
    // the accel correction fought it). Absolute, not relative to the
    // current bias, so a polluted stored value self-heals at the next true
    // still window. Dirty: only flag for persistence when the value moved
    // meaningfully (file writes, not RAM).
    double biasSaneLsb = 8.0;
    double biasDirtyLsb = 0.25;
    // Capture-stuck escape (field: LHR-599F3B91 rest bias sits beyond the
    // sane gate vs its all-zero config bias => every startup window was
    // rejected, view frozen, no console clue — even 'b' re-entered the
    // same loop). Two escapes:
    //  - 'b' (requestBiasCapture) attests stillness: the sd gate still
    //    applies, the sane gate does not — the gate only exists to catch a
    //    slow smooth rotation polluting the mean, which the user's
    //    deliberate hold-still rules out.
    //  - startup auto-escalation: biasEscalateWindows consecutive windows
    //    that pass the sd gate, fail ONLY the sane gate, and agree with
    //    each other within biasEscalateConsistLsb per axis are accepted —
    //    a worn head cannot hold a rotation constant to ~0.06°/s for
    //    seconds on end; a genuine large rest bias can.
    int biasEscalateWindows = 5;
    double biasEscalateConsistLsb = 1.0;
    // Mahony accel-correction admission (M4 step 3 corr= tightening).
    // Field finding: corr= hit 0.13-0.18 rad/s during ordinary head motion
    // — a visible tug once the neck model turns orientation into
    // translation. The |a| magnitude gate here is NARROWER than the
    // leveling gate (0.5-1.5 g, instant-start path): leveling just needs a
    // roughly-down vector, the correction term needs a CLEAN gravity
    // reference or it should stay out entirely. Rest behavior unchanged
    // (|a|=1, rate~0 => full kp).
    double corrGateMinG = 0.85;
    double corrGateMaxG = 1.15;
    double corrFadeFullRadS = 0.5;  // rate at which the fade reaches zero
    // Hard ceiling on the applied correction magnitude (rad/s). The gates
    // above bound WHEN the term acts; this bounds HOW HARD — slow head
    // translation keeps |a| inside the magnitude gate and rate inside the
    // fade while the accel direction is off-axis, and kp*fade*|e| was
    // still reaching 0.14-0.17 rad/s (field, post-tightening). At rest the
    // error is tiny and the clamp never engages; drift still pulls in.
    double corrMaxRadS = 0.10;
  };

  // cfg supplies the imu frame basis and config biases/scales.
  void configure(const HmdConfig& cfg, const Params& params);

  // Feed every unique sample, in stream order (TundraImu sample sink).
  void update(const ImuSample& s);

  // world<-head. predictSec > 0 extrapolates by the latest head-frame
  // angular velocity (q * exp(w*dt/2)) — cheap motion-to-photon latency
  // compensation (clamped to 50 ms; proper pacing/prediction is M4).
  Quat pose(double predictSec = 0.0) const;

  // Pose-age-aware prediction (M4 step 3): extrapolate from the LAST
  // SAMPLE's host arrival time to hostTimeNs (steady_clock epoch, ns) —
  // the caller passes the predicted photon time and the prediction span
  // becomes the true sample->photon age instead of a fixed guess. Spans
  // clamp to [0, 50 ms]. outAgeSec (optional) reports the clamped span
  // actually used (status-line truth). Falls back to the unpredicted pose
  // before the first sample.
  Quat poseAtHostNs(int64_t hostTimeNs, double* outAgeSec = nullptr) const;

  void recenter();         // zero yaw, keep gravity alignment
  // Redo the still-window gyro bias capture. Explicit request = user
  // attests stillness: the sane gate is waived for this capture (sd gate
  // still enforced) — see Params escalation comment.
  void requestBiasCapture();

  // Instant start: install a persisted bias (raw LSB). If the filter is
  // waiting in the startup capture, it switches to level-from-next-gravity-
  // sample and starts tracking immediately.
  void seedBias(const double biasLsb[3]);

  // True once per newly accepted bias worth persisting (startup capture,
  // 'b' recapture, or a background still-window refresh that moved the
  // value by > biasDirtyLsb); copies the current bias to outLsb. The app
  // owns throttling the actual file writes.
  bool takeBiasRefresh(double outLsb[3]);

  // Startup/explicit-capture telemetry (single slot, latest wins). The
  // silent reject loop read as a hang in the field — the app prints WHY a
  // window was rejected and when an escape path accepted. Background
  // refresh windows never emit events (rejections there are routine: any
  // head motion rejects them all day).
  struct BiasCaptureEvent {
    enum class Kind {
      kRejected,           // window evaluated and thrown away
      kAcceptedEscalated,  // consistent-streak escape accepted the window
      kAcceptedRelaxed,    // 'b' attested stillness; sane gate waived
    };
    Kind kind = Kind::kRejected;
    double mean[3] = {};  // window gyro mean, raw LSB
    double sd[3] = {};    // window gyro stddev, raw LSB
    bool sdFail = false;    // stddev over biasStillLsb (motion/vibration)
    bool saneFail = false;  // mean beyond biasSaneLsb of every sane ref
    int streak = 0;         // consecutive consistent sane-only rejects
  };
  bool takeBiasCaptureEvent(BiasCaptureEvent* out);

  struct Status {
    bool initialized = false;     // bias captured, filter running
    double gravityHead[3] = {};   // measured accel dir in head frame (unit)
    double accelMagG = 0.0;       // |accel| in g (1.0 at rest)
    double omegaRadS = 0.0;       // latest |gyro| rad/s, pre-correction —
                                  // the idle-policy motion signal (M3 step 4)
    double gyroBiasLsb[3] = {};   // captured bias
    uint64_t samples = 0;
    uint64_t resets = 0;          // hard resets (full re-level + bias capture)
    uint64_t softRecoveries = 0;  // micro-gap skips (orientation + bias kept)
    uint64_t biasRefreshes = 0;   // background still-window bias updates
  };
  Status status() const;

  // Sample tap (M4 step 4 neck cal): called once per integrated sample
  // with the head-frame RAW rate (bias-corrected gyro, no Mahony term
  // mixed in — the kinematics solve wants physical rotation), accel (g,
  // head frame), the current orientation, and dt. Runs under the filter
  // lock — keep it append-cheap. Null disarms. Tracking continues
  // unaffected while a tap is set (capture must never block tracking).
  using SampleTap = std::function<void(const double omegaHead[3],
                                       const double accHeadG[3],
                                       const Quat& q, double dt)>;
  void setSampleTap(SampleTap tap);

  // Max accel-correction magnitude applied since the last call (rad/s),
  // then resets. Rubber-band triage: with delivery, present, and reset
  // counters all clean, the remaining suspect is the Mahony tug — at slow
  // sustained head rates the fade leaves partial gain while centripetal
  // accel corrupts the gravity reference; the neck model amplifies the
  // resulting orientation drag into translation. Quiet head: ~0. Spikes
  // correlating with episodes convict this term.
  double takeMaxCorrectionRadS();

 private:
  void resetLocked(bool forceCapture);  // mu_ held
  void levelFromGravityLocked();        // mu_ held; uses st_.gravityHead
  // Sane-gate check for one axis (mu_ held): within biasSaneLsb of the
  // config rest bias, OR of the currently trusted bias when one exists.
  // The either-ref widening keeps the self-heal property (a true still
  // window near the config ref always accepts) while removing the dead
  // zone on units whose genuine rest bias exceeds the gate.
  bool saneOkLocked(double meanLsb, int axis) const;

  Params p_;
  double rHeadImu_[3][3] = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}};
  double accBias_[3] = {};   // config, g units
  double accScaleCfg_[3] = {1, 1, 1};
  double gyroBiasCfg_[3] = {};  // config, LSB

  mutable std::mutex mu_;
  Quat q_;                    // world<-head
  bool haveTc_ = false;
  uint32_t lastTc_ = 0;
  // Bias capture state. The window accumulators serve double duty: the
  // blocking startup capture, then the opportunistic background refresh.
  bool biasCapturing_ = true;
  bool haveBias_ = false;     // gyroBias_ holds a usable value (seed/capture)
  bool needLevel_ = false;    // bias known, orientation pending first gravity
  bool biasDirty_ = false;    // new bias since last takeBiasRefresh()
  bool relaxSane_ = false;    // 'b' attested stillness for this capture
  int saneRejStreak_ = 0;     // consecutive consistent sane-only rejects
  double saneRejMean_[3] = {};
  bool haveCaptureEvent_ = false;
  BiasCaptureEvent captureEvent_{};
  double biasSum_[3] = {};
  double biasSqSum_[3] = {};
  uint64_t biasN_ = 0;
  double biasTimeAcc_ = 0.0;
  double gyroBias_[3] = {};   // captured, LSB
  double lastOmega_[3] = {};  // head-frame rad/s, latest sample (prediction)
  int64_t lastSampleHostNs_ = 0;  // steady_clock arrival of latest sample
  double maxCorr_ = 0.0;      // takeMaxCorrectionRadS watermark
  SampleTap tap_;             // neck-cal feed (null = disarmed)
  Status st_;
};

}  // namespace sauna