// for MPU6050 / MPU6500 by tatsuaki kuroda KurodaENT
// This file is based on the MPU6050 library by Electronic Cats.
//Copyright (c) 2015-2023 Electronic Cats
//Licensed under the MIT License.
//https://opensource.org/licenses/MIT
//
#include <Wire.h>
#include <EEPROM.h>
#include <math.h>

// WHO_AM_IでMPU6050/MPU6500を自動判別する。
// 初回起動時だけキャリブレーションし、結果をLeonardoのEEPROMへ保存する。
// 2回目以降は保存値を読み込み、PCから'r'を受信するとTeapotデータ送信を開始する。
// 動作中または開始待ち中に小文字の'k'を受信すると、再キャリブレーションを行う。
// 再測定したオフセット値はEEPROMへ上書き保存され、その後再び'r'の受信待ちになる。
// 小文字の'c'はYaw基準の再設定であり、再キャリブレーションには使用しない。
struct SavedCalibration {
  uint32_t magic;
  uint8_t whoAmI;
  int16_t offset[6];
};

const uint32_t CALIBRATION_MAGIC = 0x6560CA2BUL;
SavedCalibration savedCalibration;
int16_t activeOffsets[6];
uint8_t detectedWhoAmI = 0;
uint8_t mpuAddress = 0x68;
// ============================================================
// =======================  CALIBRATION MODE  =================
// =======================   (code#1 run)   ===================
// ============================================================

// ====== I2C アドレス ======
// ====== レジスタ ======
#define WHO_AM_I        0x75
#define PWR_MGMT_1      0x6B
#define SMPLRT_DIV      0x19
#define CONFIG_REG      0x1A
#define GYRO_CFG        0x1B
#define ACCEL_CFG       0x1C
#define ACCEL_XOUT_H    0x3B

// Gyro user offset（共通）
#define XG_OFFS_USR_H   0x13   // H/L = 0x13/0x14
#define YG_OFFS_USR_H   0x15   //       0x15/0x16
#define ZG_OFFS_USR_H   0x17   //       0x17/0x18

// Accel offset（MPU-6500系列：H/L = 0x77/0x78, 0x7A/0x7B, 0x7D/0x7E）
// ※ LSBのbit0は予約ビット：必ず現状値のbit0を保護して書く必要あり
#define XA_OFFS_H_6500  0x77
#define YA_OFFS_H_6500  0x7A
#define ZA_OFFS_H_6500  0x7D

// MPU6050 Accel offset
#define XA_OFFS_H_6050  0x06   // H/L = 0x06/0x07
#define YA_OFFS_H_6050  0x08   //       0x08/0x09
#define ZA_OFFS_H_6050  0x0A   //       0x0A/0x0B

// ====== IMU_Zero 互換の調整パラメータ ======
const int usDelay = 3150;       // ≒200Hz相当の待ち
const int NFast   = 300;//500;//1000;       // 速い平均回数
const int NSlow   = 1500;//3000;//10000;      // 遅い平均回数
const int LinesBetweenHeaders = 5;

// 軸の添字
enum { iAx=0, iAy, iAz, iGx, iGy, iGz };

// 作業バッファ
int LowValue[6];
int HighValue[6];
int Smoothed[6];
int LowOffset[6];
int HighOffset[6];
int Target[6];
int LinesOut;
int N;
int i;

// ==== I2C helpers ====
static inline bool calWrite8(uint8_t reg, uint8_t val){
  Wire.beginTransmission(mpuAddress);
  Wire.write(reg); Wire.write(val);
  return (Wire.endTransmission()==0);
}
static inline bool calReadN(uint8_t reg, uint8_t *buf, uint8_t n){
  Wire.beginTransmission(mpuAddress);
  Wire.write(reg);
  if (Wire.endTransmission(false)!=0) return false;
  Wire.requestFrom((int)mpuAddress, (int)n);
  for(uint8_t k=0;k<n;k++){ if(!Wire.available()) return false; buf[k]=Wire.read(); }
  return true;
}
static inline uint8_t calRead8(uint8_t reg){
  Wire.beginTransmission(mpuAddress);
  Wire.write(reg);
  if (Wire.endTransmission(false)!=0) return 0xFF;
  Wire.requestFrom((int)mpuAddress, 1);
  return Wire.available()? Wire.read():0xFF;
}

bool detectMpu(){
  const uint8_t addresses[2] = {0x68, 0x69};
  for (uint8_t n=0; n<2; n++) {
    mpuAddress = addresses[n];
    uint8_t who = calRead8(WHO_AM_I);
    if (who == 0x68 || who == 0x70) {
      detectedWhoAmI = who;
      Serial.print(F("Detected "));
      Serial.print(who == 0x68 ? F("MPU6050") : F("MPU6500"));
      Serial.print(F(" at I2C address 0x"));
      Serial.println(mpuAddress, HEX);
      return true;
    }
  }
  detectedWhoAmI = 0;
  return false;
}

static inline bool calWrite16(uint8_t regH, int16_t v){
  return calWrite8(regH,   (v>>8)&0xFF)
      && calWrite8(regH+1, (v   )&0xFF);
}

// ==== Accel offset 書き込み（両機種ともLSB bit0を保護） ====
bool calWriteAccelOffsetRegister(uint8_t regH, int16_t val){
  // 既存ローバイトを読み、bit0を保護
  Wire.beginTransmission(mpuAddress);
  Wire.write(regH+1);
  if (Wire.endTransmission(false)!=0) return false;
  Wire.requestFrom((int)mpuAddress, 1);
  if (!Wire.available()) return false;
  uint8_t loOld = Wire.read();

  uint8_t hi = (val>>8) & 0xFF;
  uint8_t lo = (val    ) & 0xFF;
  lo = (uint8_t)((lo & 0xFE) | (loOld & 0x01));  // bit0保護

  return calWrite8(regH, hi) && calWrite8(regH+1, lo);
}

// 検出した機種に対応するAccel offsetレジスタを選択
bool calWriteAccelOffset(int axis, int16_t val){
  uint8_t regH;
  if (detectedWhoAmI == 0x70) {
    regH = (axis==0)? XA_OFFS_H_6500 : (axis==1)? YA_OFFS_H_6500 : ZA_OFFS_H_6500;
  } else {
    regH = (axis==0)? XA_OFFS_H_6050 : (axis==1)? YA_OFFS_H_6050 : ZA_OFFS_H_6050;
  }
  return calWriteAccelOffsetRegister(regH, val);
}

// Gyro offset 書き込み（共通）
bool calWriteGyroOffset(uint8_t regH, int16_t val){
  return calWrite16(regH, val);
}

// ==== 現オフセットセット（IMU_Zero互換）====
void SetOffsets(int TheOffsets[6]){
  // Accel
  calWriteAccelOffset(0, TheOffsets[iAx]);
  calWriteAccelOffset(1, TheOffsets[iAy]);
  calWriteAccelOffset(2, TheOffsets[iAz]);
  // Gyro
  calWriteGyroOffset(XG_OFFS_USR_H, TheOffsets[iGx]);
  calWriteGyroOffset(YG_OFFS_USR_H, TheOffsets[iGy]);
  calWriteGyroOffset(ZG_OFFS_USR_H, TheOffsets[iGz]);
}

// ==== スムージング：現在の生データ平均（オフセット込み）====
void GetSmoothed(){
  int16_t RawValue[6];
  long Sums[6]; for(i=0;i<=5;i++) Sums[i]=0;

  for (i=1; i<=N; i++){
    uint8_t raw[14];
    if (!calReadN(ACCEL_XOUT_H, raw, 14)){
      delayMicroseconds(usDelay);
      continue;
    }
    RawValue[iAx] = (int16_t)((raw[0]<<8)|raw[1]);
    RawValue[iAy] = (int16_t)((raw[2]<<8)|raw[3]);
    RawValue[iAz] = (int16_t)((raw[4]<<8)|raw[5]);
    RawValue[iGx] = (int16_t)((raw[8]<<8)|raw[9]);
    RawValue[iGy] = (int16_t)((raw[10]<<8)|raw[11]);
    RawValue[iGz] = (int16_t)((raw[12]<<8)|raw[13]);

    for (int j=0; j<=5; j++) Sums[j]+=RawValue[j];
    delayMicroseconds(usDelay);
  }
  for (i=0; i<=5; i++) Smoothed[i] = (int)((Sums[i] + N/2) / N);
}

void ForceHeader(){ LinesOut = 99; }

void ShowProgress(){
  if (LinesOut >= LinesBetweenHeaders){
    Serial.println("\t\tXAccel\t\t\tYAccel\t\t\t\tZAccel\t\t\tXGyro\t\t\tYGyro\t\t\tZGyro");
    LinesOut = 0;
  }
  Serial.print(' ');
  for (i=0;i<=5;i++){
    Serial.print('['); Serial.print(LowOffset[i]); Serial.print(',');
    Serial.print(HighOffset[i]); Serial.print("] --> [");
    Serial.print(LowValue[i]); Serial.print(','); Serial.print(HighValue[i]);
    if (i==5) Serial.println("]");
    else      Serial.print("]\t");
  }
  LinesOut++;
}

void SetAveraging(int NewN){
  N = NewN;
  Serial.print("\nAveraging "); Serial.print(N); Serial.println(" readings each time");
}

// ==== 「外へ広げる」フェーズ ====
void PullBracketsOut(){
  bool Done=false;
  int NextLowOffset[6];
  int NextHighOffset[6];

  Serial.println("Expanding:");
  ForceHeader();

  while(!Done){
    Done = true;

    // 低い方
    SetOffsets(LowOffset);
    GetSmoothed();
    for (i=0;i<=5;i++){
      LowValue[i]=Smoothed[i];
      if (LowValue[i] >= Target[i]){
        Done=false;
        NextLowOffset[i] = LowOffset[i] - 1000;
      }else{
        NextLowOffset[i] = LowOffset[i];
      }
    }

    // 高い方
    SetOffsets(HighOffset);
    GetSmoothed();
    for (i=0;i<=5;i++){
      HighValue[i]=Smoothed[i];
      if (HighValue[i] <= Target[i]){
        Done=false;
        NextHighOffset[i] = HighOffset[i] + 1000;
      }else{
        NextHighOffset[i] = HighOffset[i];
      }
    }

    ShowProgress();
    for (int k=0;k<=5;k++){ LowOffset[k]=NextLowOffset[k]; HighOffset[k]=NextHighOffset[k]; }
  }
}

// ==== 「中へ絞る」フェーズ（二分探索）====
void PullBracketsIn(){
  bool AllBracketsNarrow=false;
  bool StillWorking=true;
  int NewOffset[6];

  Serial.println("\nClosing in:");
  ForceHeader();

  while (StillWorking){
    StillWorking=false;
    if (AllBracketsNarrow && (N==NFast)) SetAveraging(NSlow);
    else AllBracketsNarrow=true;

    for (int k=0;k<=5;k++){
      if (HighOffset[k] <= (LowOffset[k]+1)){
        NewOffset[k]=LowOffset[k];
      }else{
        StillWorking=true;
        NewOffset[k]=(LowOffset[k]+HighOffset[k])/2;
        if (HighOffset[k] > (LowOffset[k]+10)) AllBracketsNarrow=false;
      }
    }

    SetOffsets(NewOffset);
    GetSmoothed();

    for (i=0;i<=5;i++){
      if (Smoothed[i] > Target[i]){
        HighOffset[i]=NewOffset[i];
        HighValue[i]=Smoothed[i];
      }else{
        LowOffset[i]=NewOffset[i];
        LowValue[i]=Smoothed[i];
      }
    }
    ShowProgress();
  }
}

// ==== 初期化 ====
void Initialize(){
  Wire.begin();
  Wire.setClock(400000);       // 不安定なら 100kHz

  Serial.begin(115200);
  delay(50);

  // 基本初期化（DLPF=1, 200Hz）
  calWrite8(PWR_MGMT_1, 0x80); delay(100);
  calWrite8(PWR_MGMT_1, 0x01); delay(10);
  calWrite8(CONFIG_REG,  0x01);
  calWrite8(SMPLRT_DIV,  0x04);   // 1kHz/(1+4)=200Hz
  calWrite8(GYRO_CFG,    0x00);   // ±250 dps
  calWrite8(ACCEL_CFG,   0x00);   // ±2 g

  // 接続確認
  Serial.println("Initializing MPU...");
  uint8_t who = calRead8(WHO_AM_I);
  Serial.print("WHO_AM_I=0x"); Serial.println(who, HEX);
  if (!(who==0x70 || who==0x68)){
    Serial.println("Unexpected WHO_AM_I. Check wiring/address/IC type.");
    while(1){}
  }
  Serial.println("MPU initialized.\n");

  Serial.println("PID tuning Each Dot = 100 readings");
  Serial.println("\nXAccel\t\tYAccel\t\tZAccel\t\tXGyro\t\tYGyro\t\tZGyro");
}

void runCalibrationAndSave(){
  Initialize();

  // 目標（Z=+1g。他の軸は0）
  for (i=iAx; i<=iGz; i++){ Target[i]=0; HighOffset[i]=0; LowOffset[i]=0; }
  Target[iAz]=16384;

  SetAveraging(NFast);
  PullBracketsOut();
  PullBracketsIn();

  // 最終オフセット（LowOffset側が解）
  Serial.println("-------------- DONE --------------");
  Serial.print  ("Final Offsets:  AX="); Serial.print(LowOffset[iAx]);
  Serial.print(" AY="); Serial.print(LowOffset[iAy]);
  Serial.print(" AZ="); Serial.print(LowOffset[iAz]);
  Serial.print("  GX="); Serial.print(LowOffset[iGx]);
  Serial.print(" GY="); Serial.print(LowOffset[iGy]);
  Serial.print(" GZ="); Serial.println(LowOffset[iGz]);

  savedCalibration.magic = CALIBRATION_MAGIC;
  savedCalibration.whoAmI = detectedWhoAmI;
  for (int k=0; k<6; k++) {
    savedCalibration.offset[k] = (int16_t)LowOffset[k];
    activeOffsets[k] = savedCalibration.offset[k];
  }
  EEPROM.put(0, savedCalibration);
  Serial.println("Offsets saved to EEPROM. Next startup will skip calibration.");
  delay(500);
}

// ============================================================
// =======================  QUATERNION MODE  ==================
// =======================   (code#2 run)   ===================
// ============================================================

// ----- Registers -----
static const uint8_t PWR_MGMT_1_Q   = 0x6B;
static const uint8_t SMPLRT_DIV_Q   = 0x19;
static const uint8_t CONFIG_REG_Q   = 0x1A;
static const uint8_t GYRO_CFG_Q     = 0x1B;
static const uint8_t ACCEL_CFG_Q    = 0x1C;
static const uint8_t ACCEL_XOUT_H_Q = 0x3B;

// ----- IO / rates -----
#define SERIAL_BAUD   230400
#define I2C_CLOCK_HZ  400000
#define DLPF_CFG_VAL  1
#define SMPL_DIV_VAL  4

// USB電源スイッチ制御ピン
const int USB_SW_EN_PIN = 9;

// ----- Scaling (±2g, ±250dps) -----
const float ACC_LSB = 16384.0f;
const float GYR_LSB = 131.0f;
const float DEG2RAD = 3.14159265358979323846f/180.0f;

// ----- Teapot 14B packet -----
uint8_t teapotPacket[14] = { '$', 0x02, 0,0,0,0,0,0,0,0, 0x00,0x00, '\r','\n' };

// ===== Quaternion state =====
struct Quat { float w,x,y,z; };
Quat q = {1,0,0,0};
bool started = false;
static inline float quatYaw(const Quat& q);

// ----- Mahony-like gains -----
float KpX = 3.8f;
float KpY = 9.0f;
float KpZ = 0.0f;
float Ki  = 0.0f;

float exInt=0.0f, eyInt=0.0f, ezInt=0.0f;

// Zジャイロ・バイアス
float gzBias = 0.0f;

// --- Yaw-hold ---
bool   haveYawRef = false;
float  yawRef     = 0.0f;

float  K_yawHold_base = 0.45f;
float  yawErrLPF = 0.0f;
const  float YAW_ERR_ALPHA = 0.15f;
const  float YAW_DEADBAND  = 1.5f*DEG2RAD;
const  float MAX_WZ_HOLD_DPS = 40.0f;
const  float D2R = 3.14159265358979323846f/180.0f;

// 静止判定
bool    stationary = false;
uint32_t stationarySinceUs = 0;
const  float ENTER_ACC_TRUST = 0.92f;
const  float EXIT_ACC_TRUST  = 0.85f;
const  float ENTER_GYRO_DPS  = 1.0f;
const  float EXIT_GYRO_DPS   = 1.5f;
const  uint32_t DEBOUNCE_US  = 250000;

// ソフトスタート
float   yawHoldGainScale = 0.0f;
const  float YAW_HOLD_RAMP_S = 0.4f;

// ===== I2C helpers =====
bool i2cWrite8(uint8_t reg, uint8_t data){
  Wire.beginTransmission(mpuAddress); Wire.write(reg); Wire.write(data);
  return (Wire.endTransmission()==0);
}
bool i2cReadN(uint8_t reg, uint8_t*buf, uint8_t n){
  Wire.beginTransmission(mpuAddress); Wire.write(reg);
  if (Wire.endTransmission(false)!=0) return false;
  Wire.requestFrom((int)mpuAddress, (int)n);
  for(uint8_t i=0;i<n;i++){ if(!Wire.available()) return false; buf[i]=Wire.read(); }
  return true;
}
bool i2cWrite16(uint8_t regH, int16_t v){
  return i2cWrite8(regH, (v>>8)&0xFF) && i2cWrite8(regH+1, v&0xFF);
}
bool writeGyroOffset(uint8_t regH, int16_t v){ return i2cWrite16(regH, v); }
bool writeAccelOffsetRegister(uint8_t regH, int16_t v){
  Wire.beginTransmission(mpuAddress); Wire.write(regH+1);
  if (Wire.endTransmission(false)!=0) return false;
  Wire.requestFrom((int)mpuAddress, 1); if(!Wire.available()) return false;
  uint8_t loOld = Wire.read();
  uint8_t hi = (v>>8)&0xFF, lo = (v)&0xFF;
  lo = (uint8_t)((lo & 0xFE) | (loOld & 0x01));   // bit0保護
  return i2cWrite8(regH, hi) && i2cWrite8(regH+1, lo);
}
bool writeAccelOffset(int axis, int16_t v){
  uint8_t regH;
  if (detectedWhoAmI == 0x70) {
    regH = (axis==0)? XA_OFFS_H_6500 : (axis==1)? YA_OFFS_H_6500 : ZA_OFFS_H_6500;
  } else {
    regH = (axis==0)? XA_OFFS_H_6050 : (axis==1)? YA_OFFS_H_6050 : ZA_OFFS_H_6050;
  }
  return writeAccelOffsetRegister(regH, v);
}
void applyHardwareOffsets(){
  writeGyroOffset(XG_OFFS_USR_H, activeOffsets[iGx]);
  writeGyroOffset(YG_OFFS_USR_H, activeOffsets[iGy]);
  writeGyroOffset(ZG_OFFS_USR_H, activeOffsets[iGz]);
  writeAccelOffset(0, activeOffsets[iAx]);
  writeAccelOffset(1, activeOffsets[iAy]);
  writeAccelOffset(2, activeOffsets[iAz]);
}

void resetFusionState(){
  q.w=1.0f; q.x=0.0f; q.y=0.0f; q.z=0.0f;
  started=false;
  exInt=0.0f; eyInt=0.0f; ezInt=0.0f;
  gzBias=0.0f;
  haveYawRef=false;
  yawRef=0.0f;
  yawErrLPF=0.0f;
  stationary=false;
  stationarySinceUs=0;
  yawHoldGainScale=0.0f;
}

void initializeQuaternionSensor(){
  Wire.begin();
  Wire.setClock(I2C_CLOCK_HZ);
  i2cWrite8(PWR_MGMT_1_Q, 0x80); delay(30);
  i2cWrite8(PWR_MGMT_1_Q, 0x01); delay(5);
  i2cWrite8(CONFIG_REG_Q,  DLPF_CFG_VAL);
  i2cWrite8(SMPLRT_DIV_Q,  SMPL_DIV_VAL);
  i2cWrite8(GYRO_CFG_Q,    0x00);
  i2cWrite8(ACCEL_CFG_Q,   0x00);
  applyHardwareOffsets();
}

void recalibrate(){
  started=false;
  // 小文字の'k'を受信したとき、データ送信を停止して再キャリブレーションする。
  // 結果はEEPROMへ保存し、姿勢計算の内部状態も初期化する。
  Serial.println(F("Recalibration requested. Keep the MPU horizontal and still."));
  runCalibrationAndSave();
  initializeQuaternionSensor();
  resetFusionState();
  Serial.println(F("Recalibration complete. New offsets saved to EEPROM."));
}

// 開始待ち：'r'=送信開始、'k'=再キャリブレーション、'F'/'O'=USB電源OFF/ON
void waitForStartSignal(){
  Serial.println(F("Waiting: 'r'=start, 'k'=recalibrate..."));
  while (true) {
    if (Serial.available()) {
      char c = Serial.read();
      if (c == 'r') {
        started = true;
        Serial.println(F("Start signal received!"));
        break;
      } else if (c == 'k') {
        recalibrate();
        Serial.println(F("Waiting: 'r'=start, 'k'=recalibrate..."));
      } else if (c == 'F') {
        digitalWrite(USB_SW_EN_PIN, LOW);
        Serial.println(F("USB power: OFF (F before start)"));
      } else if (c == 'O') {
        digitalWrite(USB_SW_EN_PIN, HIGH);
        Serial.println(F("USB power: ON (O before start)"));
      }
    }
  }
}

// ===== Quaternion utils =====
static inline float quatYaw(const Quat& q){ // ZYX yaw
  return atan2f(2.0f*(q.w*q.z + q.x*q.y),
                1.0f - 2.0f*(q.y*q.y + q.z*q.z));
}
static inline float wrapPi(float a){
  while (a >  M_PI) a -= 2.0f*M_PI;
  while (a < -M_PI) a += 2.0f*M_PI;
  return a;
}

// ===== Mahony-like update with robust yaw-hold =====
void mahonyUpdate(float gx, float gy, float gz,   // [rad/s]
                  float ax, float ay, float az,   // [g]
                  float dt,
                  float &wAccOut)
{
  static float axf=0, ayf=0, azf=0;
  const float ALP = 0.10f;
  axf += ALP*(ax-axf); ayf += ALP*(ay-ayf); azf += ALP*(az-azf);

  float norm = sqrtf(axf*axf + ayf*ayf + azf*azf);
  if (norm <= 1e-6f) return;
  float axn=axf/norm, ayn=ayf/norm, azn=azf/norm;

  float accErr = fabsf(norm - 1.0f);
  float wAcc = 1.0f - (accErr - 0.12f)/0.13f;
  if (wAcc < 0.0f) wAcc = 0.0f;
  if (wAcc > 1.0f) wAcc = 1.0f;
  wAccOut = wAcc;

  float gx_est = 2.0f*(q.x*q.z - q.w*q.y);
  float gy_est = 2.0f*(q.w*q.x + q.y*q.z);
  float gz_est = q.w*q.w - q.x*q.x - q.y*q.y + q.z*q.z;

  float ex = (ayn*gz_est - azn*gy_est);
  float ey = (azn*gx_est - axn*gz_est);

  float gx_dps = gx*(180.0f/M_PI);
  float gy_dps = gy*(180.0f/M_PI);
  float gz_dps = gz*(180.0f/M_PI);

  uint32_t nowUs = micros();
  bool enterCond = (wAcc >= ENTER_ACC_TRUST) &&
                   (fabsf(gx_dps) <= ENTER_GYRO_DPS) &&
                   (fabsf(gy_dps) <= ENTER_GYRO_DPS) &&
                   (fabsf(gz_dps) <= ENTER_GYRO_DPS);
  bool exitCond  = (wAcc <= EXIT_ACC_TRUST) ||
                   (fabsf(gx_dps) >= EXIT_GYRO_DPS) ||
                   (fabsf(gy_dps) >= EXIT_GYRO_DPS) ||
                   (fabsf(gz_dps) >= EXIT_GYRO_DPS);

  if (!stationary){
    if (enterCond){
      if (stationarySinceUs==0) stationarySinceUs = nowUs;
      if ((nowUs - stationarySinceUs) >= DEBOUNCE_US){
        stationary = true;
        yawHoldGainScale = 0.0f;
        yawRef = quatYaw(q);
        haveYawRef = true;
      }
    } else {
      stationarySinceUs = 0;
    }
  } else {
    if (exitCond){
      stationary = false;
      stationarySinceUs = 0;
      yawHoldGainScale = 0.0f;
    }
  }

  if (stationary){
    const float BETA = 0.008f;
    gzBias = (1.0f - BETA)*gzBias + BETA*gz;
    yawHoldGainScale += dt / YAW_HOLD_RAMP_S;
    if (yawHoldGainScale > 1.0f) yawHoldGainScale = 1.0f;
  }
  gz -= gzBias;

  float wx = gx + (KpX * wAcc) * ex + exInt;
  float wy = gy + (KpY * wAcc) * ey + eyInt;
  float wz = gz;

  if (stationary && haveYawRef){
    float yawNow = quatYaw(q);
    float yawErr = wrapPi(yawNow - yawRef);

    if (fabsf(yawErr) < YAW_DEADBAND) yawErr = 0.0f;
    yawErrLPF += YAW_ERR_ALPHA * (yawErr - yawErrLPF);

    float wz_hold = - (K_yawHold_base * yawHoldGainScale) * yawErrLPF;

    float wz_hold_max = MAX_WZ_HOLD_DPS * D2R;
    if (wz_hold >  wz_hold_max) wz_hold =  wz_hold_max;
    if (wz_hold < -wz_hold_max) wz_hold = -wz_hold_max;

    wz += wz_hold;
  } else {
    yawErrLPF = 0.0f;
  }

  float hw = -0.5f*( q.x*wx + q.y*wy + q.z*wz);
  float hx =  0.5f*( q.w*wx + q.y*wz - q.z*wy);
  float hy =  0.5f*( q.w*wy - q.x*wz + q.z*wx);
  float hz =  0.5f*( q.w*wz + q.x*wy - q.y*wx);

  q.w += hw*dt; q.x += hx*dt; q.y += hy*dt; q.z += hz*dt;

  float n = sqrtf(q.w*q.w + q.x*q.x + q.y*q.y + q.z*q.z);
  if (n>0.0f){ float inv=1.0f/n; q.w*=inv; q.x*=inv; q.y*=inv; q.z*=inv; }
}

void setup(){
  Serial.begin(SERIAL_BAUD);
  while(!Serial) {}

  Wire.begin();
  Wire.setClock(I2C_CLOCK_HZ);
  if (!detectMpu()) {
    Serial.println(F("MPU6050/MPU6500 not found. Check wiring and I2C address."));
    while (true) {}
  }

  EEPROM.get(0, savedCalibration);
  if (savedCalibration.magic != CALIBRATION_MAGIC ||
      savedCalibration.whoAmI != detectedWhoAmI) {
    Serial.println(F("No saved calibration. Starting first-time calibration."));
    Serial.println(F("Keep the MPU horizontal and completely still."));
    runCalibrationAndSave();
  } else {
    for (int k=0; k<6; k++) activeOffsets[k] = savedCalibration.offset[k];
    Serial.println(F("Calibration loaded from EEPROM."));
  }

  pinMode(USB_SW_EN_PIN, OUTPUT);
  digitalWrite(USB_SW_EN_PIN, HIGH);

  initializeQuaternionSensor();
  resetFusionState();
  waitForStartSignal();
}

void loop(){
  while (Serial.available()){
    int c = Serial.read();
    if (c == 'c'){
      yawRef = quatYaw(q);
      haveYawRef = true;
    } else if (c == 'k') {
      recalibrate();
      waitForStartSignal();
    } else if (c == 'F') {
      digitalWrite(USB_SW_EN_PIN, LOW);
      Serial.println(F("USB power: OFF (F)"));
    } else if (c == 'O') {
      digitalWrite(USB_SW_EN_PIN, HIGH);
      Serial.println(F("USB power: ON (O)"));
    }
  }

  if (!started) return;

  uint8_t raw[14];
  if (!i2cReadN(ACCEL_XOUT_H_Q, raw, 14)) return;

  int16_t ax_i=(raw[0]<<8)|raw[1];
  int16_t ay_i=(raw[2]<<8)|raw[3];
  int16_t az_i=(raw[4]<<8)|raw[5];
  int16_t gx_i=(raw[8]<<8)|raw[9];
  int16_t gy_i=(raw[10]<<8)|raw[11];
  int16_t gz_i=(raw[12]<<8)|raw[13];

  float ax = ax_i/ACC_LSB;
  float ay = ay_i/ACC_LSB;
  float az = az_i/ACC_LSB;
  float gx = (gx_i/GYR_LSB)*DEG2RAD;
  float gy = (gy_i/GYR_LSB)*DEG2RAD;
  float gz = (gz_i/GYR_LSB)*DEG2RAD;

  static uint32_t tprev = micros();
  uint32_t tnow = micros();
  float dt = (tnow - tprev)*1e-6f;
  if (dt <= 0.0f || dt > 0.02f) dt = 1.0f/200.0f;
  tprev = tnow;

  float wAcc=0.0f;
  mahonyUpdate(gx, gy, gz, ax, ay, az, dt, wAcc);

  auto packQ = [](float v)->int16_t{
    float s=v*16384.0f; if(s>32767)s=32767; if(s<-32768)s=-32768; return (int16_t)s;
  };
  int16_t q1=packQ(q.w), q2=packQ(q.x), q3=packQ(q.y), q4=packQ(q.z);

  teapotPacket[2]=(uint8_t)((q1>>8)&0xFF);
  teapotPacket[3]=(uint8_t)( q1     &0xFF);
  teapotPacket[4]=(uint8_t)((q2>>8)&0xFF);
  teapotPacket[5]=(uint8_t)( q2     &0xFF);
  teapotPacket[6]=(uint8_t)((q3>>8)&0xFF);
  teapotPacket[7]=(uint8_t)( q3     &0xFF);
  teapotPacket[8]=(uint8_t)((q4>>8)&0xFF);
  teapotPacket[9]=(uint8_t)( q4     &0xFF);
  Serial.write(teapotPacket, 14);
  teapotPacket[11]++;
}