#include <M5Atom.h>
#include <Wire.h>
#include <Preferences.h>
#include <BLEDevice.h>
#include <BLEServer.h>
#include <BLEUtils.h>
#include <BLE2902.h>
#include <math.h>

const char DEVICE_NAME[] = "ATOM-MPU6886";
const char SERVICE_UUID[] = "7f1b0001-8e7a-4d76-9f47-7b5a9e9d0001";
const char DATA_UUID[] = "7f1b0002-8e7a-4d76-9f47-7b5a9e9d0001";
const char COMMAND_UUID[] = "7f1b0003-8e7a-4d76-9f47-7b5a9e9d0001";
// ESP32-S3版と同じく、センサー読み出し・姿勢演算・BLE通知を100 Hzで行う。
// 200 Hzで融合演算しても送信は100 Hzのため、ATOMではBLE処理との競合が増える。
const uint32_t SAMPLE_INTERVAL_US = 10000;
const uint32_t BLE_INTERVAL_US = 10000;
const uint32_t CALIBRATION_MAGIC = 0x6886A701;
const uint8_t FAST_IMU_ADDRESS = 0x68;
const uint8_t FAST_IMU_ACCEL_XOUT_H = 0x3B;

BLECharacteristic *dataCharacteristic = nullptr;
Preferences preferences;
volatile bool connected = false, startRequested = false, yawResetRequested = false;
bool streaming = false;
float biasX = 0, biasY = 0, biasZ = 0;
float qw = 1, qx = 0, qy = 0, qz = 0, yawOffset = 0;
// yaw補正用の固定クォータニオン。送信のたびにcosf()/sinf()を呼ばない。
float yawCW = 1.0f, yawSZ = 0.0f;

static void updateYawRotation() {
  const float h = -yawOffset * 0.5f;
  yawCW = cosf(h);
  yawSZ = sinf(h);
}

static bool readImuFast(float &ax, float &ay, float &az,
                        float &gx, float &gy, float &gz) {
  uint8_t raw[14];
  Wire1.beginTransmission(FAST_IMU_ADDRESS);
  Wire1.write(FAST_IMU_ACCEL_XOUT_H);
  if (Wire1.endTransmission(false) != 0) return false;
  if (Wire1.requestFrom((int)FAST_IMU_ADDRESS, (int)sizeof(raw)) != sizeof(raw)) return false;
  for (uint8_t i = 0; i < sizeof(raw); ++i) {
    if (!Wire1.available()) return false;
    raw[i] = Wire1.read();
  }
  const int16_t axRaw = (int16_t)((raw[0] << 8) | raw[1]);
  const int16_t ayRaw = (int16_t)((raw[2] << 8) | raw[3]);
  const int16_t azRaw = (int16_t)((raw[4] << 8) | raw[5]);
  const int16_t gxRaw = (int16_t)((raw[8] << 8) | raw[9]);
  const int16_t gyRaw = (int16_t)((raw[10] << 8) | raw[11]);
  const int16_t gzRaw = (int16_t)((raw[12] << 8) | raw[13]);
  ax = axRaw * (2.0f / 32768.0f); ay = ayRaw * (2.0f / 32768.0f); az = azRaw * (2.0f / 32768.0f);
  gx = gxRaw * (250.0f / 32768.0f); gy = gyRaw * (250.0f / 32768.0f); gz = gzRaw * (250.0f / 32768.0f);
  return true;
}

class ServerCallbacks : public BLEServerCallbacks {
  void onConnect(BLEServer *) override { connected = true; }
  void onDisconnect(BLEServer *) override {
    connected = false; streaming = false; BLEDevice::startAdvertising();
  }
};

class CommandCallbacks : public BLECharacteristicCallbacks {
  void onWrite(BLECharacteristic *c) override {
    String value = c->getValue();
    for (size_t i = 0; i < value.length(); ++i) {
      if (value[i] == 'r') startRequested = true;
      if (value[i] == 'c') yawResetRequested = true;
    }
  }
};

void calibrateOnce() {
  float b[3];
  if (preferences.getUInt("magic", 0) == CALIBRATION_MAGIC &&
      preferences.getBytesLength("bias") == sizeof(b)) {
    preferences.getBytes("bias", b, sizeof(b));
    biasX = b[0]; biasY = b[1]; biasZ = b[2];
    Serial.println("Gyro calibration loaded");
    return;
  }

  Serial.println("Keep ATOM still: calibrating for 5 seconds...");
  double sx = 0, sy = 0, sz = 0;
  for (int i = 0; i < 1000; ++i) {
    float x, y, z;
    M5.IMU.getGyroData(&x, &y, &z);
    sx += x; sy += y; sz += z;
    delay(5);
  }
  biasX = sx / 1000.0; biasY = sy / 1000.0; biasZ = sz / 1000.0;
  b[0] = biasX; b[1] = biasY; b[2] = biasZ;
  preferences.putBytes("bias", b, sizeof(b));
  preferences.putUInt("magic", CALIBRATION_MAGIC);
  Serial.println("Gyro calibration saved");
}

void updateFusion(float gx, float gy, float gz, float ax, float ay, float az, float dt) {
  // ESP32-S3版と同じMahony系処理：加速度LPFと信頼度重み付けを使用する。
  static float axf = 0.0f, ayf = 0.0f, azf = 0.0f;
  const float alpha = 0.10f;
  axf += alpha * (ax - axf);
  ayf += alpha * (ay - ayf);
  azf += alpha * (az - azf);

  float n = sqrtf(axf * axf + ayf * ayf + azf * azf);
  if (n < 0.01f) return;
  const float accError = fabsf(n - 1.0f);
  float accWeight = 1.0f - (accError - 0.12f) / 0.13f;
  accWeight = constrain(accWeight, 0.0f, 1.0f);
  ax = axf / n; ay = ayf / n; az = azf / n;
  float vx = 2 * (qx * qz - qw * qy);
  float vy = 2 * (qw * qx + qy * qz);
  float vz = qw * qw - qx * qx - qy * qy + qz * qz;
  // Mahony型の加速度補正。加速度からヨー(Z軸)は補正しない。
  float ex = ay * vz - az * vy;
  float ey = az * vx - ax * vz;
  gx += 3.8f * accWeight * ex;
  gy += 9.0f * accWeight * ey;

  float h = 0.5f * dt;
  float dw = (-qx * gx - qy * gy - qz * gz) * h;
  float dx = ( qw * gx + qy * gz - qz * gy) * h;
  float dy = ( qw * gy - qx * gz + qz * gx) * h;
  float dz = ( qw * gz + qx * gy - qy * gx) * h;
  qw += dw; qx += dx; qy += dy; qz += dz;
  n = sqrtf(qw*qw + qx*qx + qy*qy + qz*qz);
  qw /= n; qx /= n; qy /= n; qz /= n;
}

int16_t q14(float v) { return (int16_t)lroundf(constrain(v, -1.0f, 1.0f) * 16384.0f); }

void setup() {
  M5.begin(true, true, true);
  M5.IMU.Init();
  Wire1.setClock(400000UL);
  M5.IMU.SetGyroFsr(MPU6886::GFS_250DPS);
  M5.IMU.SetAccelFsr(MPU6886::AFS_2G);
  preferences.begin("atom-imu", false);
  calibrateOnce();

  BLEDevice::init(DEVICE_NAME);
  BLEServer *server = BLEDevice::createServer();
  server->setCallbacks(new ServerCallbacks());
  BLEService *service = server->createService(SERVICE_UUID);
  dataCharacteristic = service->createCharacteristic(DATA_UUID, BLECharacteristic::PROPERTY_READ | BLECharacteristic::PROPERTY_NOTIFY);
  dataCharacteristic->addDescriptor(new BLE2902());
  BLECharacteristic *command = service->createCharacteristic(COMMAND_UUID, BLECharacteristic::PROPERTY_WRITE | BLECharacteristic::PROPERTY_WRITE_NR);
  command->setCallbacks(new CommandCallbacks());
  service->start();
  BLEAdvertising *advertising = BLEDevice::getAdvertising();
  advertising->addServiceUUID(SERVICE_UUID);
  advertising->setScanResponse(true);
  BLEDevice::startAdvertising();
}

void loop() {
  M5.update();
  if (startRequested) { startRequested = false; streaming = true; }
  if (yawResetRequested) {
    yawResetRequested = false;
    yawOffset = atan2f(2 * (qw*qz + qx*qy), 1 - 2 * (qy*qy + qz*qz));
    updateYawRotation();
  }
  if (!connected || !streaming) { delay(1); return; }

  static uint32_t nextUs = 0, lastBleUs = 0;
  uint32_t now = micros();
  if (!nextUs) nextUs = now;
  if ((int32_t)(now - nextUs) < 0) return;
  nextUs += SAMPLE_INTERVAL_US;
  if ((int32_t)(now - nextUs) > 50000) nextUs = now;

  static uint32_t previousUs = 0;
  float dt = previousUs ? (now - previousUs) * 1e-6f : SAMPLE_INTERVAL_US * 1e-6f;
  previousUs = now;
  if (dt <= 0.0f || dt > 0.02f) dt = SAMPLE_INTERVAL_US * 1e-6f;

  float ax, ay, az, gx, gy, gz;
  if (!readImuFast(ax, ay, az, gx, gy, gz)) return;
  updateFusion((gx-biasX)*DEG_TO_RAD, (gy-biasY)*DEG_TO_RAD, (gz-biasZ)*DEG_TO_RAD, ax, ay, az, dt);
  now = micros();
  if ((uint32_t)(now - lastBleUs) < BLE_INTERVAL_US) return;
  lastBleUs = now;

  const float cw = yawCW, sz = yawSZ;
  const float w = cw*qw - sz*qz, x = cw*qx + sz*qy;
  const float y = cw*qy - sz*qx, z = cw*qz + sz*qw;
  uint8_t packet[14] = {'$', 0x02, 0,0,0,0,0,0,0,0,0,0,'\r','\n'};
  int16_t v[4] = {q14(w), q14(x), q14(y), q14(z)};
  for (int i = 0; i < 4; ++i) { packet[2+i*2] = v[i] >> 8; packet[3+i*2] = v[i]; }
  dataCharacteristic->setValue(packet, sizeof(packet));
  dataCharacteristic->notify();
}