#include "imu.h"
#include "motion.h"
#include "Wire.h"
#include <math.h>

LSM6DS3 imu(I2C_MODE, 0x6A);

uint32_t imuFifoOverruns = 0;

// Registers not in config.h
#define REG_FIFO_CTRL3   0x08
#define REG_FIFO_CTRL5   0x0A
#define REG_CTRL2_G      0x11
#define REG_CTRL3_C      0x12
#define REG_CTRL4_C      0x13
#define REG_FIFO_STATUS1 0x3A
#define REG_FIFO_DATA_L  0x3E

#define FIFO_CTRL5_CONT_208HZ 0x2E  // ODR_FIFO=208 Hz, FIFO_MODE=continuous

static const int CAL_TIMEOUT_MS = 3000;
// Gyro peak-to-peak (any axis) above this during calibration means the device moved.
static const float CAL_STILL_PP_DPS = 1.0f;

// I2C helpers
void imuWriteReg(uint8_t reg, uint8_t val) {
  // LSM6DS3 is on Wire1 (internal I2C, SDA=P0.17, SCL=P0.16), NOT Wire (external pins 4/5)
  Wire1.beginTransmission(0x6A); Wire1.write(reg); Wire1.write(val); Wire1.endTransmission();
}

uint8_t imuReadReg(uint8_t reg) {
  Wire1.beginTransmission(0x6A); Wire1.write(reg); Wire1.endTransmission(false);
  Wire1.requestFrom((uint8_t)0x6A, (uint8_t)1);
  return Wire1.available() ? Wire1.read() : 0;
}

// Burst read using Wire1's TWIM directly, with the hardware generating the repeated
// start and the STOP (LASTTX->STARTRX, LASTRX->STOP shortcuts).
// Wire1.requestFrom() triggers STOP from software after the LASTRX event; when a
// SoftDevice interrupt delays that, the TWIM clocks in extra bytes. Here that is
// not harmless: an over-read of FIFO_STATUS runs into FIFO_DATA_OUT and silently
// consumes FIFO words, which shows up as corrupted, shifted samples.
uint32_t imuI2cErrors = 0;

static bool imuReadBlock(uint8_t reg, uint8_t* buf, uint16_t len) {
  NRF_TWIM_Type* t = NRF_TWIM1;
  static uint8_t txReg;
  txReg = reg;
  t->EVENTS_STOPPED = 0; t->EVENTS_ERROR = 0; t->EVENTS_SUSPENDED = 0;
  t->EVENTS_LASTTX = 0;  t->EVENTS_LASTRX = 0;
  t->EVENTS_TXSTARTED = 0; t->EVENTS_RXSTARTED = 0;
  t->ERRORSRC = t->ERRORSRC;
  t->ADDRESS = 0x6A;
  t->TXD.PTR = (uint32_t)&txReg; t->TXD.MAXCNT = 1;
  t->RXD.PTR = (uint32_t)buf;    t->RXD.MAXCNT = len;
  t->SHORTS = TWIM_SHORTS_LASTTX_STARTRX_Msk | TWIM_SHORTS_LASTRX_STOP_Msk;
  t->TASKS_RESUME = 1;
  t->TASKS_STARTTX = 1;
  uint32_t t0 = micros();
  bool ok = true;
  while (!t->EVENTS_STOPPED) {
    if (t->EVENTS_ERROR || micros() - t0 > 20000) {
      ok = false;
      t->TASKS_STOP = 1;
      uint32_t t1 = micros();
      while (!t->EVENTS_STOPPED && micros() - t1 < 2000) {}
      break;
    }
  }
  t->SHORTS = 0;
  if (t->RXD.AMOUNT != len) ok = false;
  t->EVENTS_STOPPED = 0; t->EVENTS_ERROR = 0;
  t->EVENTS_LASTTX = 0;  t->EVENTS_LASTRX = 0;
  t->EVENTS_TXSTARTED = 0; t->EVENTS_RXSTARTED = 0;
  t->ERRORSRC = t->ERRORSRC;
  if (!ok) imuI2cErrors++;
  return ok;
}

// Configuration
// Called after imu.begin(). The library settings are kept in sync in main.cpp so
// its readFloat*() helpers (used by the sleep manager) still scale correctly.
void imuConfigure() {
  Wire1.setClock(400000);
  imuWriteReg(REG_CTRL3_C,  0x44);  // BDU=1, IF_INC=1
  imuWriteReg(REG_CTRL4_C,  0x00);  // clear bit7 set by the library (DEN_XL_EN on the TR-C)
  imuWriteReg(REG_CTRL1_XL, 0x58);  // accel 208 Hz, ±4 g
  imuWriteReg(REG_CTRL2_G,  0x58);  // gyro  208 Hz, ±1000 dps
  imuWriteReg(REG_FIFO_CTRL3, 0x09);  // gyro + accel into FIFO, no decimation
  delay(150);  // gyro turn-on / reconfiguration transient (first samples read 100+ dps off)
  imuFifoReset();
}

void imuFifoReset() {
  imuWriteReg(REG_FIFO_CTRL5, 0x00);  // bypass mode empties the FIFO
  imuWriteReg(REG_FIFO_CTRL5, FIFO_CTRL5_CONT_208HZ);
}

// Drain up to maxSamples complete gyro+accel sets from the FIFO.
// Each set is 6 words in the order Gx Gy Gz Ax Ay Az. FIFO_PATTERN gives the
// index of the next word, so a partial set (after an overrun) is skipped first.
int imuReadFifo(ImuRaw* out, int maxSamples) {
  uint8_t st[4];
  if (!imuReadBlock(REG_FIFO_STATUS1, st, 4)) return 0;
  uint16_t words   = st[0] | ((uint16_t)(st[1] & 0x07) << 8);
  uint16_t pattern = st[2] | ((uint16_t)(st[3] & 0x03) << 8);
  if (st[1] & 0x40) imuFifoOverruns++;

  while (pattern != 0 && words > 0) {
    uint8_t skip[2];
    if (!imuReadBlock(REG_FIFO_DATA_L, skip, 2)) return 0;
    words--; pattern = (pattern + 1) % 6;
  }

  int n = words / 6;
  if (n > maxSamples) n = maxSamples;
  const int PER_READ = 16;
  int done = 0;
  while (done < n) {
    int chunk = min(PER_READ, n - done);
    uint8_t buf[PER_READ * 12];
    // Multi-byte reads of FIFO_DATA_OUT roll back from _H to _L, so a burst
    // returns consecutive FIFO words.
    if (!imuReadBlock(REG_FIFO_DATA_L, buf, chunk * 12)) break;
    for (int i = 0; i < chunk; i++) {
      const uint8_t* b = buf + i * 12;
      ImuRaw& s = out[done + i];
      s.gx = (int16_t)(b[0]  | (b[1]  << 8));
      s.gy = (int16_t)(b[2]  | (b[3]  << 8));
      s.gz = (int16_t)(b[4]  | (b[5]  << 8));
      s.ax = (int16_t)(b[6]  | (b[7]  << 8));
      s.ay = (int16_t)(b[8]  | (b[9]  << 8));
      s.az = (int16_t)(b[10] | (b[11] << 8));
    }
    done += chunk;
  }
  return done;
}

// Temperature
float readIMUTemp() {
  int16_t raw = (int16_t)((imuReadReg(REG_OUT_TEMP_H) << 8) | imuReadReg(REG_OUT_TEMP_L));
  return 25.0f + (float)raw / 256.0f;
}

// Calibration
// Averages ~1 s of FIFO samples. If the device moved, the result is only used
// when there is no previous bias; the continuous bias tracker refines it later.
void calibrateGyroBias() {
  Serial.println("[CAL] Hold still...");
  imuFifoReset();
  double sum[6] = {0};
  int16_t mn[3] = {32767, 32767, 32767}, mx[3] = {-32768, -32768, -32768};
  int n = 0;
  unsigned long t0 = millis();
  ImuRaw buf[16];
  while (n < BIAS_SAMPLES && millis() - t0 < (unsigned long)CAL_TIMEOUT_MS) {
    int k = imuReadFifo(buf, 16);
    for (int i = 0; i < k; i++) {
      const int16_t v[6] = {buf[i].gx, buf[i].gy, buf[i].gz, buf[i].ax, buf[i].ay, buf[i].az};
      for (int j = 0; j < 6; j++) sum[j] += v[j];
      for (int j = 0; j < 3; j++) { mn[j] = min(mn[j], v[j]); mx[j] = max(mx[j], v[j]); }
    }
    n += k;
    digitalWrite(LED_GREEN, ((millis() - t0) / 100) % 2);  // green flutter during calibration
    delay(5);
  }
  digitalWrite(LED_GREEN, HIGH);

  if (n == 0) { Serial.println("[CAL] No IMU samples - calibration skipped"); return; }

  float pp = 0;
  for (int j = 0; j < 3; j++) pp = max(pp, (mx[j] - mn[j]) * GYRO_DPS_PER_LSB);
  bool still = pp < CAL_STILL_PP_DPS;
  float bias[3], accel[3];
  for (int j = 0; j < 3; j++) {
    bias[j]  = (float)(sum[j] / n) * GYRO_RAD_PER_LSB;
    accel[j] = (float)(sum[j + 3] / n) * ACCEL_G_PER_LSB;
  }
  calTempC = readIMUTemp();
  bool applied = motionApplyCalibration(bias, accel, still);

  #ifdef FEATURE_TELEMETRY
    statRecalCount++;
  #endif

  Serial.print("[CAL] T="); Serial.print(calTempC, 1);
  Serial.print("C n="); Serial.print(n);
  Serial.print(" p-p="); Serial.print(pp, 2); Serial.print("dps");
  Serial.print(" bias(dps)="); Serial.print(bias[0] * (180.f / PI), 3);
  Serial.print(","); Serial.print(bias[1] * (180.f / PI), 3);
  Serial.print(","); Serial.print(bias[2] * (180.f / PI), 3);
  Serial.println(applied ? (still ? " applied" : " applied (moved - will refine when still)")
                         : " ignored (device moved, keeping previous bias)");
  imuFifoReset();
}

// Motion curve
float applyCurve(float v) {
  switch (cfg.curve) {
    case CURVE_SQUARE: return (v >= 0 ? 1.f : -1.f) * v * v;
    case CURVE_SQRT:   return (v >= 0 ? 1.f : -1.f) * sqrtf(fabsf(v));
    default:           return v;
  }
}

float applyAcceleration(float d) { return d * (1.0f + fabsf(d) * cfg.accelStrength); }