#include "RemoteIMU.h"
#include "Config.h"
#include <Wire.h>

// MPU6050 register map, default +/-2g / +/-250deg/s ranges (same part and
// scale factors as robot/RobotIMU.cpp).
static constexpr uint8_t REG_CONFIG       = 0x1A;
static constexpr uint8_t REG_PWR_MGMT_1   = 0x6B;
static constexpr uint8_t REG_ACCEL_XOUT_H = 0x3B;

static bool writeReg(uint8_t reg, uint8_t val) {
  IMU_WIRE.beginTransmission(IMU_ADDR);
  IMU_WIRE.write(reg);
  IMU_WIRE.write(val);
  return IMU_WIRE.endTransmission() == 0;
}

bool RemoteIMU::begin() {
  IMU_WIRE.begin();
  IMU_WIRE.setClock(400000);
  ok_ = writeReg(REG_PWR_MGMT_1, 0x00);   // wake out of sleep
  if (!ok_) return false;
  writeReg(REG_CONFIG, 0x03);             // 44 Hz DLPF: a handheld remote doesn't need more bandwidth
  delay(50);

  // Gyro bias: average while the remote sits still at power-on.
  float sum[3] = {};
  int got = 0;
  for (int n = 0; n < IMU_GYRO_BIAS_SAMPLES; n++) {
    float a[3], g[3];
    if (readRaw(a, g)) {
      for (int k = 0; k < 3; k++) sum[k] += g[k];
      got++;
    }
    delay(2);
  }
  if (got > 0)
    for (int k = 0; k < 3; k++) gyro_bias_[k] = sum[k] / got;

  // Seed the filter from the accelerometer so it doesn't slew in from 0,
  // and treat the power-on attitude as level until the operator re-zeroes.
  float a[3], g[3];
  if (readRaw(a, g)) {
    roll_deg_  = degrees(atan2f(a[1], a[2]));
    pitch_deg_ = degrees(atan2f(-a[0], sqrtf(a[1] * a[1] + a[2] * a[2])));
  }
  zeroLevel();
  return true;
}

bool RemoteIMU::readRaw(float a[3], float g[3]) {
  IMU_WIRE.beginTransmission(IMU_ADDR);
  IMU_WIRE.write(REG_ACCEL_XOUT_H);
  if (IMU_WIRE.endTransmission(false) != 0) return false;
  IMU_WIRE.requestFrom((int)IMU_ADDR, 14);
  if (IMU_WIRE.available() < 14) return false;

  int16_t raw[7];   // ax ay az temp gx gy gz, big-endian
  for (int k = 0; k < 7; k++) {
    uint8_t hi = IMU_WIRE.read();   // two statements: the order of calls inside one expression is unspecified
    uint8_t lo = IMU_WIRE.read();
    raw[k] = (int16_t)((hi << 8) | lo);
  }

  float chip_a[3] = { raw[0] / 16384.0f, raw[1] / 16384.0f, raw[2] / 16384.0f };
  float chip_g[3] = { raw[4] / 131.0f,   raw[5] / 131.0f,   raw[6] / 131.0f };

  for (int k = 0; k < 3; k++) {
    a[k] = chip_a[IMU_AXIS_FROM_CHIP[k]] * IMU_AXIS_SIGN[k];
    g[k] = chip_g[IMU_AXIS_FROM_CHIP[k]] * IMU_AXIS_SIGN[k] - gyro_bias_[k];
  }
  return true;
}

void RemoteIMU::update(float dt) {
  if (!ok_) return;
  float a[3], g[3];
  if (!readRaw(a, g)) return;

  float accel_roll  = degrees(atan2f(a[1], a[2]));
  float accel_pitch = degrees(atan2f(-a[0], sqrtf(a[1] * a[1] + a[2] * a[2])));

  roll_deg_  = IMU_ALPHA * (roll_deg_  + g[0] * dt) + (1 - IMU_ALPHA) * accel_roll;
  pitch_deg_ = IMU_ALPHA * (pitch_deg_ + g[1] * dt) + (1 - IMU_ALPHA) * accel_pitch;
}