#include <stdint.h>
#include <stdio.h>

#include "bsp/mmio.h"
#include "bsp/pmod_spi.h"
#include "bsp/soc.h"
#include "bsp/uart.h"

/* ICM-42686-P/ICM-42688-P user bank 0 registers. */
#define ICM_DEVICE_CONFIG     0x11u
#define ICM_TEMP_DATA1        0x1du
#define ICM_INT_STATUS        0x2du
#define ICM_PWR_MGMT0         0x4eu
#define ICM_GYRO_CONFIG0      0x4fu
#define ICM_ACCEL_CONFIG0     0x50u
#define ICM_GYRO_CONFIG1      0x51u
#define ICM_GYRO_ACCEL_CONFIG0 0x52u
#define ICM_ACCEL_CONFIG1     0x53u
#define ICM_INT_SOURCE0       0x65u
#define ICM_WHO_AM_I          0x75u

#define ICM_SPI_READ          0x80u
#define ICM_SPI_DUMMY         0xffu
#define ICM_DATA_READY        0x08u

#define ICM42686_WHO_AM_I     0x44u
#define ICM42688_WHO_AM_I     0x47u

#define ICM_ODR_500HZ         0x0fu
#define ICM_PWR_LOW_NOISE     0x0fu
#define ICM_FILTER_BW_12P5HZ  0x77u

#define CALIBRATION_SAMPLES   2500u
#define CALIBRATION_DISCARD   100u
#define CALIBRATION_ATTEMPTS  3u
#define CALIBRATION_TIMEOUT_S 8u
#define REPORT_PERIOD_MS      500u
#ifndef REPORT_DURATION_S
#define REPORT_DURATION_S     60u
#endif
#define REPORT_COUNT          ((REPORT_DURATION_S * 1000u) / REPORT_PERIOD_MS)

#define IMU_AXIS_COUNT        6u
#define IMU_FIELD_COUNT       7u

typedef enum {
  CAPTURE_IDLE = 0,
  CAPTURE_CALIBRATION,
  CAPTURE_RUN
} capture_mode_t;

typedef struct {
  int16_t temperature;
  int16_t accel[3];
  int16_t gyro[3];
} imu_raw_sample_t;

static volatile capture_mode_t capture_mode = CAPTURE_IDLE;
static volatile uint32_t calibration_discard;
static volatile uint32_t calibration_count;
static volatile int64_t calibration_sum[IMU_AXIS_COUNT];
static volatile uint64_t calibration_sum_sq[IMU_AXIS_COUNT];

static volatile uint32_t run_window_count;
static volatile uint32_t run_total_count;
static volatile int64_t run_window_sum[IMU_FIELD_COUNT];

static volatile uint32_t invalid_sample_count;

static int32_t accel_bias_raw[3];
static int32_t gyro_bias_raw[3];
static int32_t gyro_full_scale_mdps;

static uint64_t time_now(void)
{
  return bsp_read64(BSP_CLINT_MTIME);
}

static uint64_t ticks_from_ms(uint32_t ms)
{
  return ((uint64_t)BSP_CLINT_TIMEBASE_HZ * ms + 999u) / 1000u;
}

static void delay_us(uint32_t us)
{
  const uint64_t ticks =
      ((uint64_t)BSP_CLINT_TIMEBASE_HZ * us + 999999u) / 1000000u;
  const uint64_t start = time_now();
  while ((time_now() - start) < ticks) {
  }
}

static void delay_ms(uint32_t ms)
{
  const uint64_t ticks = ticks_from_ms(ms);
  const uint64_t start = time_now();
  while ((time_now() - start) < ticks) {
  }
}

static void icm_write_reg(uint8_t reg, uint8_t value)
{
  bsp_pmod_spi_write_reg8(reg & (uint8_t)~ICM_SPI_READ, value);
}

static uint8_t icm_read_reg(uint8_t reg)
{
  return bsp_pmod_spi_read_reg8(reg | ICM_SPI_READ, ICM_SPI_DUMMY);
}

static int16_t be16(const uint8_t *p)
{
  return (int16_t)(((uint16_t)p[0] << 8) | p[1]);
}

static void icm_read_sample(imu_raw_sample_t *sample)
{
  uint8_t raw[14];

  bsp_pmod_spi_read_regs(ICM_TEMP_DATA1 | ICM_SPI_READ,
                         ICM_SPI_DUMMY,
                         raw,
                         sizeof(raw));
  sample->temperature = be16(&raw[0]);
  sample->accel[0] = be16(&raw[2]);
  sample->accel[1] = be16(&raw[4]);
  sample->accel[2] = be16(&raw[6]);
  sample->gyro[0] = be16(&raw[8]);
  sample->gyro[1] = be16(&raw[10]);
  sample->gyro[2] = be16(&raw[12]);
}

static int sample_valid(const imu_raw_sample_t *sample)
{
  if (sample->temperature == INT16_MIN) {
    return 0;
  }

  for (unsigned i = 0; i < 3u; ++i) {
    if ((sample->accel[i] == INT16_MIN) ||
        (sample->gyro[i] == INT16_MIN)) {
      return 0;
    }
  }
  return 1;
}

static void capture_sample(const imu_raw_sample_t *sample)
{
  if (!sample_valid(sample)) {
    invalid_sample_count++;
    return;
  }

  if (capture_mode == CAPTURE_CALIBRATION) {
    if (calibration_discard != 0u) {
      calibration_discard--;
      return;
    }

    if (calibration_count < CALIBRATION_SAMPLES) {
      const int32_t values[IMU_AXIS_COUNT] = {
        sample->accel[0], sample->accel[1], sample->accel[2],
        sample->gyro[0], sample->gyro[1], sample->gyro[2]
      };

      for (unsigned i = 0; i < IMU_AXIS_COUNT; ++i) {
        calibration_sum[i] += values[i];
        calibration_sum_sq[i] += (uint64_t)((int64_t)values[i] * values[i]);
      }
      calibration_count++;
    }
    return;
  }

  if (capture_mode == CAPTURE_RUN) {
    run_window_sum[0] += sample->temperature;
    run_window_sum[1] += sample->accel[0];
    run_window_sum[2] += sample->accel[1];
    run_window_sum[3] += sample->accel[2];
    run_window_sum[4] += sample->gyro[0];
    run_window_sum[5] += sample->gyro[1];
    run_window_sum[6] += sample->gyro[2];
    run_window_count++;
    run_total_count++;
  }
}

static int poll_data_ready(void)
{
  if ((icm_read_reg(ICM_INT_STATUS) & ICM_DATA_READY) == 0u) {
    return 0;
  }

  imu_raw_sample_t sample;
  icm_read_sample(&sample);
  capture_sample(&sample);
  return 1;
}

static int64_t divide_round_nearest(int64_t numerator, int64_t denominator)
{
  if (numerator >= 0) {
    return (numerator + denominator / 2) / denominator;
  }
  return -((-numerator + denominator / 2) / denominator);
}

static uint64_t integer_sqrt(uint64_t value)
{
  uint64_t result = 0u;
  uint64_t bit = UINT64_C(1) << 62;

  while (bit > value) {
    bit >>= 2;
  }
  while (bit != 0u) {
    if (value >= result + bit) {
      value -= result + bit;
      result = (result >> 1) + bit;
    } else {
      result >>= 1;
    }
    bit >>= 2;
  }
  return result;
}

static void reset_calibration_capture(void)
{
  capture_mode = CAPTURE_IDLE;
  calibration_discard = CALIBRATION_DISCARD;
  calibration_count = 0u;
  for (unsigned i = 0; i < IMU_AXIS_COUNT; ++i) {
    calibration_sum[i] = 0;
    calibration_sum_sq[i] = 0u;
  }
  capture_mode = CAPTURE_CALIBRATION;
}

static int wait_for_calibration(void)
{
  const uint64_t timeout_ticks =
      (uint64_t)BSP_CLINT_TIMEBASE_HZ * CALIBRATION_TIMEOUT_S;
  const uint64_t start = time_now();

  while (calibration_count < CALIBRATION_SAMPLES) {
    (void)poll_data_ready();
    if ((time_now() - start) >= timeout_ticks) {
      return 0;
    }
  }
  return 1;
}

static int calculate_calibration(void)
{
  int64_t sum[IMU_AXIS_COUNT];
  uint64_t sum_sq[IMU_AXIS_COUNT];
  int64_t mean[IMU_AXIS_COUNT];
  uint64_t sigma[IMU_AXIS_COUNT];
  const int64_t accel_counts_per_g = 16384;

  capture_mode = CAPTURE_IDLE;
  for (unsigned i = 0; i < IMU_AXIS_COUNT; ++i) {
    sum[i] = calibration_sum[i];
    sum_sq[i] = calibration_sum_sq[i];
  }
  for (unsigned i = 0; i < IMU_AXIS_COUNT; ++i) {
    mean[i] = divide_round_nearest(sum[i], CALIBRATION_SAMPLES);
    const uint64_t mean_square = sum_sq[i] / CALIBRATION_SAMPLES;
    const uint64_t square_of_mean = (uint64_t)(mean[i] * mean[i]);
    const uint64_t variance =
        mean_square > square_of_mean ? mean_square - square_of_mean : 0u;
    sigma[i] = integer_sqrt(variance);
  }

  const uint64_t accel_norm = integer_sqrt(
      (uint64_t)(mean[0] * mean[0]) +
      (uint64_t)(mean[1] * mean[1]) +
      (uint64_t)(mean[2] * mean[2]));
  const uint64_t accel_sigma_limit = (uint64_t)accel_counts_per_g / 50u;
  const uint64_t gyro_sigma_limit =
      (uint64_t)divide_round_nearest((int64_t)32768 * 1000,
                                    gyro_full_scale_mdps);
  const int64_t gyro_mean_limit =
      divide_round_nearest((int64_t)32768 * 2000, gyro_full_scale_mdps);

  int stable =
      (accel_norm >= (uint64_t)(accel_counts_per_g * 8 / 10)) &&
      (accel_norm <= (uint64_t)(accel_counts_per_g * 12 / 10));

  for (unsigned i = 0; i < 3u; ++i) {
    if ((sigma[i] > accel_sigma_limit) ||
        (sigma[i + 3u] > gyro_sigma_limit) ||
        (mean[i + 3u] < -gyro_mean_limit) ||
        (mean[i + 3u] > gyro_mean_limit)) {
      stable = 0;
    }
  }

  printf("cal raw mean A=%ld,%ld,%ld G=%ld,%ld,%ld norm=%lu\n",
         (long)mean[0], (long)mean[1], (long)mean[2],
         (long)mean[3], (long)mean[4], (long)mean[5],
         (unsigned long)accel_norm);
  printf("cal raw sigma A=%lu,%lu,%lu G=%lu,%lu,%lu: %s\n",
         (unsigned long)sigma[0], (unsigned long)sigma[1],
         (unsigned long)sigma[2], (unsigned long)sigma[3],
         (unsigned long)sigma[4], (unsigned long)sigma[5],
         stable ? "PASS" : "MOVING/UNSTABLE");

  if (!stable || (accel_norm == 0u)) {
    return 0;
  }

  for (unsigned i = 0; i < 3u; ++i) {
    const int64_t expected_gravity = divide_round_nearest(
        mean[i] * accel_counts_per_g, (int64_t)accel_norm);
    accel_bias_raw[i] = (int32_t)(mean[i] - expected_gravity);
    gyro_bias_raw[i] = (int32_t)mean[i + 3u];
  }

  printf("cal bias raw A=%ld,%ld,%ld G=%ld,%ld,%ld\n",
         (long)accel_bias_raw[0], (long)accel_bias_raw[1],
         (long)accel_bias_raw[2], (long)gyro_bias_raw[0],
         (long)gyro_bias_raw[1], (long)gyro_bias_raw[2]);
  return 1;
}

static void reset_run_capture(void)
{
  capture_mode = CAPTURE_IDLE;
  run_window_count = 0u;
  run_total_count = 0u;
  for (unsigned i = 0; i < IMU_FIELD_COUNT; ++i) {
    run_window_sum[i] = 0;
  }
  capture_mode = CAPTURE_RUN;
}

static uint32_t take_run_window(int64_t sum[IMU_FIELD_COUNT],
                                uint32_t *total_count)
{
  const uint32_t count = run_window_count;

  for (unsigned i = 0; i < IMU_FIELD_COUNT; ++i) {
    sum[i] = run_window_sum[i];
    run_window_sum[i] = 0;
  }
  run_window_count = 0u;
  *total_count = run_total_count;
  return count;
}

static void print_run_window(unsigned report_index)
{
  int64_t sum[IMU_FIELD_COUNT];
  uint32_t total_count;
  const uint32_t count = take_run_window(sum, &total_count);

  if (count == 0u) {
    printf("%02u.%03us n=0 total=%u NO POLLING DATA\n",
           (report_index * REPORT_PERIOD_MS) / 1000u,
           (report_index * REPORT_PERIOD_MS) % 1000u,
           total_count);
    return;
  }

  double mean[IMU_FIELD_COUNT];
  for (unsigned i = 0; i < IMU_FIELD_COUNT; ++i) {
    mean[i] = (double)sum[i] / (double)count;
  }

  const double temperature_c = 25.0 + mean[0] / 132.48;
  const double ax_g = (mean[1] - (double)accel_bias_raw[0]) / 16384.0;
  const double ay_g = (mean[2] - (double)accel_bias_raw[1]) / 16384.0;
  const double az_g = (mean[3] - (double)accel_bias_raw[2]) / 16384.0;
  const double gyro_scale_dps = (double)gyro_full_scale_mdps / 1000.0;
  const double gx_dps =
      (mean[4] - (double)gyro_bias_raw[0]) * gyro_scale_dps / 32768.0;
  const double gy_dps =
      (mean[5] - (double)gyro_bias_raw[1]) * gyro_scale_dps / 32768.0;
  const double gz_dps =
      (mean[6] - (double)gyro_bias_raw[2]) * gyro_scale_dps / 32768.0;

  printf("%02u.%03u s | samples=%3u total=%5u | "
         "Temp=%6.3f C | "
         "Accel[g] X=% .3f Y=% .3f Z=% .3f | "
         "Gyro[dps] X=% .3f Y=% .3f Z=% .3f\n",
         (report_index * REPORT_PERIOD_MS) / 1000u,
         (report_index * REPORT_PERIOD_MS) % 1000u,
         count, total_count, temperature_c,
         ax_g, ay_g, az_g, gx_dps, gy_dps, gz_dps);
}

int main(void)
{
  uint8_t gyro_config0;
  uint8_t accel_config0;

  printf("\nICM PMOD 500 Hz SPI polling test\n");
  printf("SPI=0x%08lx Genesys2 PMOD JC\n",
         (unsigned long)BSP_PMOD_SPI_BASE);
  printf("SPI pins: CS=JC1 MOSI=JC2 MISO=JC3 SCK=JC4\n");
  printf("INT1/JC9 is not used; samples are read by polling INT_STATUS.\n");

  bsp_pmod_spi_init();
  bsp_pmod_spi_set_mode(BSP_PMOD_SPI_MODE_SINGLE);
  bsp_pmod_spi_set_clock_mode(BSP_PMOD_SPI_CLOCK_MODE0);
  bsp_pmod_spi_set_timing(10u, 10u, 10u, 10u);
  delay_us(1000u);

  const uint8_t who_am_i = icm_read_reg(ICM_WHO_AM_I);
  const uint8_t device_config = icm_read_reg(ICM_DEVICE_CONFIG);
  if (who_am_i == ICM42688_WHO_AM_I) {
    gyro_config0 = (7u << 5) | ICM_ODR_500HZ;
    accel_config0 = (3u << 5) | ICM_ODR_500HZ;
    gyro_full_scale_mdps = 15625;
    printf("WHO_AM_I=0x%02x ICM-42688-P\n", who_am_i);
  } else if (who_am_i == ICM42686_WHO_AM_I) {
    gyro_config0 = (7u << 5) | ICM_ODR_500HZ;
    accel_config0 = (4u << 5) | ICM_ODR_500HZ;
    gyro_full_scale_mdps = 31250;
    printf("WHO_AM_I=0x%02x ICM-42686-P\n", who_am_i);
  } else {
    printf("WHO_AM_I=0x%02x FAIL (expected 0x44 or 0x47)\n", who_am_i);
    bsp_uart_wait_tx_idle();
    return 1;
  }

  if (device_config != 0x00u) {
    printf("DEVICE_CONFIG=0x%02x FAIL (SPI mode must remain 0)\n",
           device_config);
    bsp_uart_wait_tx_idle();
    return 1;
  }

  /* Configuration registers are changed while accel and gyro are off. */
  delay_ms(50u);
  icm_write_reg(ICM_PWR_MGMT0, 0x00u);
  delay_ms(1u);
  icm_write_reg(ICM_INT_SOURCE0, 0x00u);

  icm_write_reg(ICM_GYRO_CONFIG0, gyro_config0);
  icm_write_reg(ICM_ACCEL_CONFIG0, accel_config0);

  /* Third-order UI/DEC2 filters and 12.5 Hz LPF at a 500 Hz ODR. */
  const uint8_t gyro_config1 =
      (icm_read_reg(ICM_GYRO_CONFIG1) & 0x10u) | 0xcau;
  const uint8_t accel_config1 =
      (icm_read_reg(ICM_ACCEL_CONFIG1) & 0xe1u) | 0x14u;
  icm_write_reg(ICM_GYRO_CONFIG1, gyro_config1);
  icm_write_reg(ICM_ACCEL_CONFIG1, accel_config1);
  icm_write_reg(ICM_GYRO_ACCEL_CONFIG0, ICM_FILTER_BW_12P5HZ);

  /* Do not route data-ready to INT1; this test uses SPI polling only. */
  icm_write_reg(ICM_INT_SOURCE0, 0x00u);

  icm_write_reg(ICM_PWR_MGMT0, ICM_PWR_LOW_NOISE);
  delay_ms(50u);

  printf("CFG PWR=%02x GYRO=%02x ACCEL=%02x FILT=%02x G1=%02x A1=%02x\n",
         icm_read_reg(ICM_PWR_MGMT0), icm_read_reg(ICM_GYRO_CONFIG0),
         icm_read_reg(ICM_ACCEL_CONFIG0),
         icm_read_reg(ICM_GYRO_ACCEL_CONFIG0),
         icm_read_reg(ICM_GYRO_CONFIG1), icm_read_reg(ICM_ACCEL_CONFIG1));
  printf("POLL INT_SOURCE0=%02x (expected 00: INT1 disabled)\n",
         icm_read_reg(ICM_INT_SOURCE0));

  if ((icm_read_reg(ICM_PWR_MGMT0) != ICM_PWR_LOW_NOISE) ||
      (icm_read_reg(ICM_GYRO_CONFIG0) != gyro_config0) ||
      (icm_read_reg(ICM_ACCEL_CONFIG0) != accel_config0) ||
      (icm_read_reg(ICM_GYRO_ACCEL_CONFIG0) != ICM_FILTER_BW_12P5HZ) ||
      (icm_read_reg(ICM_INT_SOURCE0) != 0x00u)) {
    printf("Sensor configuration read-back failed\n");
    bsp_uart_wait_tx_idle();
    return 1;
  }

  /* Clear any stale data-ready status before polling starts. */
  (void)icm_read_reg(ICM_INT_STATUS);
  capture_mode = CAPTURE_IDLE;

  printf("Keep the board motionless for calibration; starting in 3 seconds.\n");
  delay_ms(3000u);

  int calibrated = 0;
  for (unsigned attempt = 1u;
       (attempt <= CALIBRATION_ATTEMPTS) && !calibrated;
       ++attempt) {
    printf("Calibration %u/%u: collecting %u samples (5 s)...\n",
           attempt, CALIBRATION_ATTEMPTS, CALIBRATION_SAMPLES);
    reset_calibration_capture();
    if (!wait_for_calibration()) {
      printf("Calibration timeout: got %u/%u polling samples.\n",
             calibration_count, CALIBRATION_SAMPLES);
      const uint8_t status = icm_read_reg(ICM_INT_STATUS);
      printf("DIAG INT_STATUS=%02x PWR=%02x\n", status,
             icm_read_reg(ICM_PWR_MGMT0));
      break;
    }
    calibrated = calculate_calibration();
  }

  if (!calibrated) {
    capture_mode = CAPTURE_IDLE;
    icm_write_reg(ICM_PWR_MGMT0, 0x00u);
    delay_ms(1u);
    icm_write_reg(ICM_INT_SOURCE0, 0x00u);
    printf("Calibration failed; keep the kit still and restart the test.\n");
    bsp_uart_wait_tx_idle();
    return 1;
  }

  printf("Starting %u s capture: 500 Hz SPI polling, UART every 500 ms.\n",
         (unsigned)REPORT_DURATION_S);
  printf("Each line is the average of samples read in that 500 ms window.\n");
  reset_run_capture();

  const uint64_t run_start = time_now();
  const uint64_t report_ticks = ticks_from_ms(REPORT_PERIOD_MS);
  for (unsigned report = 1u; report <= REPORT_COUNT; ++report) {
    const uint64_t deadline = run_start + (uint64_t)report * report_ticks;
    while (time_now() < deadline) {
      (void)poll_data_ready();
    }
    print_run_window(report);
  }

  capture_mode = CAPTURE_IDLE;
  icm_write_reg(ICM_PWR_MGMT0, 0x00u);
  delay_ms(1u);
  icm_write_reg(ICM_INT_SOURCE0, 0x00u);
  (void)icm_read_reg(ICM_INT_STATUS);

  printf("Complete: samples=%u expected<=%u invalid=%u\n",
         run_total_count, 500u * REPORT_DURATION_S,
         invalid_sample_count);
  bsp_uart_wait_tx_idle();
  return 0;
}