#include "ugv_io_mpu.hh"

#include <driver/uart.h>
#include <esp_log.h>

#include <cmath>

constexpr float M_PI = 3.1415926535897932384626433832795;

#include "MPU.hpp"
#include "i2c_mutex.h"
#include "mpu/math.hpp"

namespace ugv {
namespace io {

static constexpr gpio_num_t MPU_SDA = GPIO_NUM_5;
static constexpr gpio_num_t MPU_SCL = GPIO_NUM_4;
static constexpr mpud::accel_fs_t MPU_ACCEL_FS = mpud::ACCEL_FS_2G;
static constexpr mpud::gyro_fs_t MPU_GYRO_FS = mpud::GYRO_FS_500DPS;
static constexpr float MPU_MAG_TO_FLUX = (4912.f) / (32760.f);

static const Vec3f ACCEL_OFFSET = {0., 0., 0.};
static const Mat3f ACCEL_MAT = {1., 0., 0., 0., 1., 0., 0., 0., 1.};
//static const Vec3f MAG_OFFSET   = {76.6938, 11.5453, 64.3988};
//static const Vec3f MAG_OFFSET   = {62.3746, 9.89597,39.9587};
static const Vec3f MAG_OFFSET   = {70.8462, 19.7919, 64.0239};
//static const Vec3f MAG_OFFSET = {0, 0, 0,};
//static const Mat3f MAG_MAT = {0., -0.0305935, 0., 0.0294454, 0.,
//                              0., 0.,         0., 0.0289344};
//  static const Mat3f MAG_MAT = {0., -0.034202, 0.,0.028624, 0., 0., 0., 0., 0.0276165};
//static const Mat3f MAG_MAT = {0., -0.028624, 0.,0.034202, 0., 0., 0., 0., 0.0276165};
static const Mat3f MAG_MAT = {0., -0.0279054, 0., 0.0272776, 0., 0., 0., 0., 0.0259509};
//static const Mat3f MAG_MAT = {1, 0, 0, 0, 1, 0, 0, 0, 1};
//static const Vec3f GYRO_OFFSET = {-2.01854, -2.62769, -0.30177};
static const Vec3f GYRO_OFFSET = {-1.99247, -2.48463, 0.294875};
//static const Vec3f GYRO_OFFSET = {0, 0, 0};
static const Mat3f GYRO_MAT = {1., 0., 0., 0., 1., 0., 0., 0., -1.};

static const char *TAG = "ugv_io_mpu";

Vec3f::Vec3f() : x(0), y(0), z(0) {}
Vec3f::Vec3f(float x, float y, float z) : x(x), y(y), z(z) {}
Vec3f::Vec3f(const mpud::float_axes_t &axes)
    : x(axes.x), y(axes.y), z(axes.z) {}

MPU::MPU() : mpu_(nullptr) {}

MPU::~MPU() { delete mpu_; }

void MPU::Init() {
  mutex_ = xSemaphoreCreateMutex();

  BaseType_t xRet = xTaskCreate(MPU::TaskEntry, "ugv_io_mpu", 8 * 1024, this, 3,
                                &this->task_);
  if (xRet != pdTRUE) {
    ESP_LOGE(TAG, "error creating MPU task");
    return;
  }
}

void MPU::Calibrate() {
  mpud::raw_axes_t accel_offset, gyro_offset;
  mpu_->computeOffsets(&accel_offset, &gyro_offset);
  mpu_->setAccelOffset(accel_offset);
  mpu_->setGyroOffset(gyro_offset);
  {
    auto ao = mpud::math::accelGravity(accel_offset, mpud::ACCEL_FS_2G);
    auto go = mpud::math::gyroDegPerSec(gyro_offset, mpud::GYRO_FS_250DPS);
    ESP_LOGI(TAG, "MPU offsets: accel=(%f, %f, %f) gyro=(%f, %f, %f)", ao.x,
             ao.y, ao.z, go.x, go.y, go.z);
  }
}

void MPU::GetData(MpuData &data) {
  xSemaphoreTake(mutex_, pdMS_TO_TICKS(10));
  data = last_data_;
  last_data_.gyro_rate = {0, 0, 0};
  xSemaphoreGive(mutex_);
}

bool MPU::Initialize() {
  xSemaphoreTake(i2c_mutex, portMAX_DELAY);
  mpu_bus_ = &i2c0;
  // This is shared with the oled, so just use those pins
  mpu_bus_->setTimeout(10);
  mpu_bus_->begin(MPU_SDA, MPU_SCL, GPIO_PULLUP_DISABLE, GPIO_PULLUP_DISABLE, 400000);
  mpu_ = new mpud::MPU(*mpu_bus_);

  esp_err_t ret;
  int tries = 0;
  for (; tries < 5; ++tries) {
    mpu_->getInterruptStatus();
    ret = mpu_->testConnection();
    if (ret != ESP_OK) {
      uint8_t wai = mpu_->whoAmI();
      ESP_LOGE(TAG, "MPU not connected (whoAmI: %#x)", wai);
      vTaskDelay(pdMS_TO_TICKS(100));
      continue;
    }
    ret = mpu_->compassTestConnection();
    if (ret != ESP_OK) {
      uint8_t wai = mpu_->compassWhoAmI();
      ESP_LOGW(TAG, "MPU compass not connected (whoAmI: %#x)", wai);
    }
    ret = mpu_->initialize();
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "MPU initialization error");
      continue;
    }
    break;
  }
  if (tries == 5) {
    return false;
  }
  // Calibrate();
  mpu_->setAccelFullScale(MPU_ACCEL_FS);
  mpu_->setGyroFullScale(MPU_GYRO_FS);
  // force magnetometer into continuous mode
  mpu_->compassWriteByte(0x0A, 0x12);
  xSemaphoreGive(i2c_mutex);

  ESP_LOGI(TAG, "MPU initialized");
  return true;
}

void MPU::DoTask() {
  ESP_LOGI(TAG, "MPU task start");
  mpud::raw_axes_t accel_, gyro_, mag_;
  esp_err_t ret;
  MpuData data;

  if (!Initialize()) {
    return;
  }

  while (true) {
    xSemaphoreTake(i2c_mutex, portMAX_DELAY);
    ret = mpu_->motion(&accel_, &gyro_);
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "error reading MPU: %#x", ret);
      xSemaphoreGive(i2c_mutex);
      continue;
    }
    uint8_t compass_data[7];
    mpu_->setAuxI2CBypass(true);
    ret = mpu_bus_->readBytes(mpud::COMPASS_I2CADDRESS, 0x03, 7, compass_data);
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "error reading MPU compass: %#x", ret);
      xSemaphoreGive(i2c_mutex);
      continue;
    }
    mpu_->setAuxI2CBypass(false);
    xSemaphoreGive(i2c_mutex);
    int16_t mx = (static_cast<int16_t>(compass_data[1]) << 8) |
                 static_cast<int16_t>(compass_data[0]);
    int16_t my = (static_cast<int16_t>(compass_data[3]) << 8) |
                 static_cast<int16_t>(compass_data[2]);
    int16_t mz = (static_cast<int16_t>(compass_data[5]) << 8) |
                 static_cast<int16_t>(compass_data[4]);
    //  ESP_LOGV(TAG, "compass: %d, %d, %d", mx, my, mz);
    data.accel = mpud::accelGravity(accel_, MPU_ACCEL_FS);
    data.accel = ACCEL_MAT * (data.accel + ACCEL_OFFSET);
    data.gyro_rate = mpud::gyroDegPerSec(gyro_, MPU_GYRO_FS);
    data.gyro_rate = GYRO_MAT * (data.gyro_rate + GYRO_OFFSET);
    data.mag.x = ((float)mx) * MPU_MAG_TO_FLUX;
    data.mag.y = ((float)my) * MPU_MAG_TO_FLUX;
    data.mag.z = ((float)mz) * MPU_MAG_TO_FLUX;
    data.mag = MAG_MAT * (data.mag + MAG_OFFSET);
    xSemaphoreTake(mutex_, pdMS_TO_TICKS(100));
    last_data_ = data;
    xSemaphoreGive(mutex_);
  }
}

void MPU::TaskEntry(void *arg) {
  MPU *mpu = (MPU *)arg;
  mpu->DoTask();
  vTaskDelete(NULL);
}

};  // namespace io
};  // namespace ugv