#include "RobotIMU.h"
#include <Wire.h>

// MPU6050-style register map, default +/-2g / +/-250deg/s ranges.
// Swap readRaw() if you're using a different IMU part.
static constexpr uint8_t IMU_ADDR         = 0x68;
static constexpr uint8_t REG_PWR_MGMT_1   = 0x6B;
static constexpr uint8_t REG_ACCEL_XOUT_H = 0x3B;

bool RobotIMU::begin() {
  Wire.begin();
  Wire.beginTransmission(IMU_ADDR);
  Wire.write(REG_PWR_MGMT_1);
  Wire.write(0x00);   // wake the IMU out of sleep
  ok_ = (Wire.endTransmission() == 0);
  return ok_;
}

bool RobotIMU::readRaw(float& ax, float& ay, float& az, float& gx, float& gy, float& gz) {
  Wire.beginTransmission(IMU_ADDR);
  Wire.write(REG_ACCEL_XOUT_H);
  if (Wire.endTransmission(false) != 0) return false;
  Wire.requestFrom((int)IMU_ADDR, 14);
  if (Wire.available() < 14) return false;

  int16_t raw_ax = (Wire.read() << 8) | Wire.read();
  int16_t raw_ay = (Wire.read() << 8) | Wire.read();
  int16_t raw_az = (Wire.read() << 8) | Wire.read();
  Wire.read(); Wire.read();   // skip temperature
  int16_t raw_gx = (Wire.read() << 8) | Wire.read();
  int16_t raw_gy = (Wire.read() << 8) | Wire.read();
  int16_t raw_gz = (Wire.read() << 8) | Wire.read();

  ax = raw_ax / 16384.0f; ay = raw_ay / 16384.0f; az = raw_az / 16384.0f;
  gx = raw_gx / 131.0f;   gy = raw_gy / 131.0f;   gz = raw_gz / 131.0f;   // deg/s
  (void)gz;
  return true;
}

void RobotIMU::update(float dt) {
  if (!ok_ || !enabled) return;
  float ax, ay, az, gx, gy, gz;
  if (!readRaw(ax, ay, az, gx, gy, gz)) return;

  float accel_roll  = degrees(atan2(ay, az));
  float accel_pitch = degrees(atan2(-ax, sqrt(ay * ay + az * az)));

  constexpr float ALPHA = 0.98f;   // trust the gyro short-term, the accelerometer long-term
  roll_deg_  = ALPHA * (roll_deg_  + gx * dt) + (1 - ALPHA) * accel_roll;
  pitch_deg_ = ALPHA * (pitch_deg_ + gy * dt) + (1 - ALPHA) * accel_pitch;
}