#include "InertialSensor_Invensense.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

#define GRAVITY_MSS 9.80665f
#define AP_INERTIAL_SENSOR_ACCEL_CLIP_THRESH_MSS (15.5f * GRAVITY_MSS) // accelerometer values over 15.5G are recorded as a clipping error

#define debug(fmt, args...)               \
    do                                    \
    {                                     \
        printf("MPU: " fmt "\n", ##args); \
    } while (0)
#define info(fmt, args...)                \
    do                                    \
    {                                     \
        printf("MPU: " fmt "\n", ##args); \
    } while (0)

#ifndef MIN
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
#endif
#ifndef MAX
#define MAX(a, b) (((a) > (b)) ? (a) : (b))
#endif

/*
  EXT_SYNC allows for frame synchronisation with an external device
  such as a camera. When enabled the LSB of AccelZ holds the FSYNC bit
 */
#ifndef INVENSENSE_EXT_SYNC_ENABLE
#define INVENSENSE_EXT_SYNC_ENABLE 0
#endif

#include "InertialSensor_Invensense_registers.h"

#define MPU_SAMPLE_SIZE 14
#define MPU_FIFO_BUFFER_LEN 16

#define int16_val(v, idx) ((int16_t)(((uint16_t)v[2 * idx] << 8) | v[2 * idx + 1]))
#define uint16_val(v, idx) (((uint16_t)v[2 * idx] << 8) | v[2 * idx + 1])

// acclerometers on Invensense sensors will return values up to
// 24G, but they are not guaranteed to be remotely linear past
// 16G
static const uint16_t multiplier_accel = INT16_MAX / (26 * GRAVITY_MSS);

/* Initialize sensor*/
static bool _hardware_init(InertialSensor_Invensense_t *inven);
static void _start(InertialSensor_Invensense_t *inven);
static bool _check_whoami(InertialSensor_Invensense_t *inven);

static void _set_filter_register(InertialSensor_Invensense_t *inven);
static void _fifo_reset(InertialSensor_Invensense_t *inven);

/* Read samples from FIFO (FIFO enabled) */
static void _read_fifo(InertialSensor_Invensense_t *inven);

/* Check if there's data available by either reading DRDY pin or register */
static bool _data_ready(InertialSensor_Invensense_t *inven);

/* Read and write functions taking the differences between buses into
* account */
static bool _block_read(InertialSensor_Invensense_t *inven, uint8_t reg, uint8_t *buf, uint32_t size);
static uint8_t _register_read(InertialSensor_Invensense_t *inven, uint8_t reg);
static void _register_write(InertialSensor_Invensense_t *inven, uint8_t reg, uint8_t val);

static bool _accumulate(InertialSensor_Invensense_t *inven, uint8_t *samples, uint8_t n_samples);

static bool _check_raw_temp(InertialSensor_Invensense_t *inven, int16_t t2);

/*
 *  RM-MPU-6000A-00.pdf, page 33, section 4.25 lists LSB sensitivity of
 *  gyro as 16.4 LSB/DPS at scale factor of +/- 2000dps (FS_SEL==3)
 */
static const float GYRO_SCALE = (0.0174532f / 16.4f);

/*
 *  RM-MPU-6000A-00.pdf, page 31, section 4.23 lists LSB sensitivity of
 *  accel as 4096 LSB/mg at scale factor of +/- 8g (AFS_SEL==2)
 *
 *  See note below about accel scaling of engineering sample MPU6k
 *  variants however
 */

bool InertialSensor_Invensense_init(InertialSensor_Invensense_t *inven, const char *name, SPI_DEV_t *dev, GPIO_TypeDef *drdy, uint16_t drdy_Pin)
{
    inven->temp_sensitivity = 1.0 / 340; // degC/LSB
    inven->temp_zero = 36.53;            // degC
    inven->_dev = dev;
    inven->_drdy = drdy;
    inven->_drdy_Pin = drdy_Pin;
    inven->name = name;

    inven->accel_clip_count = 0;
    inven->_accum.count = 0;
    inven->_accum.accel[0] = 0;
    inven->_accum.accel[1] = 0;
    inven->_accum.accel[2] = 0;
    inven->_accum.gyro[0] = 0;
    inven->_accum.gyro[1] = 0;
    inven->_accum.gyro[2] = 0;

    SPI_DEV_set_read_flag(dev, 0x80);

    inven->success = _hardware_init(inven);

    if (inven->success)
    {
        _start(inven);
    }

    osMutexDef(inven);
    inven->mutex = osMutexCreate(osMutex(inven));
    if (!inven->mutex)
    {
        return false;
    }

    return inven->success;
}

void _fifo_reset(InertialSensor_Invensense_t *inven)
{
    uint8_t user_ctrl = inven->_last_stat_user_ctrl;
    user_ctrl &= ~(BIT_USER_CTRL_FIFO_RESET | BIT_USER_CTRL_FIFO_EN);

    _register_write(inven, MPUREG_FIFO_EN, 0);
    _register_write(inven, MPUREG_USER_CTRL, user_ctrl);
    _register_write(inven, MPUREG_USER_CTRL, user_ctrl | BIT_USER_CTRL_FIFO_RESET);
    _register_write(inven, MPUREG_USER_CTRL, user_ctrl | BIT_USER_CTRL_FIFO_EN);
    _register_write(inven, MPUREG_FIFO_EN, BIT_XG_FIFO_EN | BIT_YG_FIFO_EN | BIT_ZG_FIFO_EN | BIT_ACCEL_FIFO_EN | BIT_TEMP_FIFO_EN);
    vTaskDelay(1);
    inven->_last_stat_user_ctrl = user_ctrl | BIT_USER_CTRL_FIFO_EN;

    //notify_accel_fifo_reset(_accel_instance);
    //notify_gyro_fifo_reset(_gyro_instance);
}

static void _start(InertialSensor_Invensense_t *inven)
{
    if (!SPI_DEV_begin(inven->_dev, 1000))
    {
        return;
    }

    // only used for wake-up in accelerometer only low power mode
    _register_write(inven, MPUREG_PWR_MGMT_2, 0x00);
    vTaskDelay(1);

    // always use FIFO
    _fifo_reset(inven);
    inven->_accum.count = 0;

    /*
      setup temperature sensitivity and offset. This varies
      considerably between parts
     */
    switch (inven->_mpu_type)
    {
    case Invensense_MPU9250:
        inven->temp_zero = 21;
        inven->temp_sensitivity = 1.0 / 340;
        break;

    case Invensense_MPU6000:
    case Invensense_MPU6500:
        inven->temp_zero = 36.53;
        inven->temp_sensitivity = 1.0 / 340;
        break;

    case Invensense_ICM20608:
    case Invensense_ICM20602:
    case Invensense_ICM20600:
        inven->temp_zero = 25;
        inven->temp_sensitivity = 1.0 / 326.8;
        break;

    case Invensense_ICM20789:
        inven->temp_zero = 25;
        inven->temp_sensitivity = 0.003;
        break;
    case Invensense_ICM20689:
        inven->temp_zero = 25;
        inven->temp_sensitivity = 0.003;
        break;
    }

    // setup ODR and on-sensor filtering
    _set_filter_register(inven);

    // set sample rate to 1000Hz and apply a software filter
    // In this configuration, the gyro sample rate is 8kHz
    _register_write(inven, MPUREG_SMPLRT_DIV, 0);
    vTaskDelay(1);

    // Gyro scale 2000o/s
    _register_write(inven, MPUREG_GYRO_CONFIG, BITS_GYRO_FS_2000DPS);
    vTaskDelay(1);

    // read the product ID rev c has 1/2 the sensitivity of rev d
    uint8_t product_id = _register_read(inven, MPUREG_PRODUCT_ID);

    if (inven->_mpu_type == Invensense_MPU6000 &&
        ((product_id == MPU6000ES_REV_C4) ||
         (product_id == MPU6000ES_REV_C5) ||
         (product_id == MPU6000_REV_C4) ||
         (product_id == MPU6000_REV_C5)))
    {
        // Accel scale 8g (4096 LSB/g)
        // Rev C has different scaling than rev D
        _register_write(inven, MPUREG_ACCEL_CONFIG, 1 << 3);
        inven->_accel_scale = GRAVITY_MSS / 4096.f;
    }
    else
    {
        // Accel scale 16g (2048 LSB/g)
        _register_write(inven, MPUREG_ACCEL_CONFIG, 3 << 3);
        inven->_accel_scale = GRAVITY_MSS / 2048.f;
    }
    vTaskDelay(1);

    if (inven->_mpu_type == Invensense_ICM20608 ||
        inven->_mpu_type == Invensense_ICM20602 ||
        inven->_mpu_type == Invensense_ICM20600)
    {
        // this avoids a sensor bug, see description above
        _register_write(inven, MPUREG_ICM_UNDOC1, MPUREG_ICM_UNDOC1_VALUE);
    }

    // configure interrupt to fire when new data arrives
    _register_write(inven, MPUREG_INT_ENABLE, BIT_RAW_RDY_EN);
    vTaskDelay(1);

    // clear interrupt on any read, and hold the data ready pin high
    // until we clear the interrupt. We don't do this for the 20789 as
    // that sensor has already setup the appropriate config inside the
    // baro driver.
    if (inven->_mpu_type != Invensense_ICM20789)
    {
        uint8_t v = _register_read(inven, MPUREG_INT_PIN_CFG) | BIT_INT_RD_CLEAR | BIT_LATCH_INT_EN;
        v &= BIT_BYPASS_EN;
        _register_write(inven, MPUREG_INT_PIN_CFG, v);
    }

    SPI_DEV_end(inven->_dev);

    // allocate fifo buffer
    inven->_fifo_buffer = (uint8_t *)pvPortMalloc(MPU_FIFO_BUFFER_LEN * MPU_SAMPLE_SIZE);
    if (inven->_fifo_buffer == NULL)
    {
        info("[%s] Unable to allocate FIFO buffer", inven->name);
        Error_Handler();
    }
}

/*
 * Return true if the Invensense has new data available for reading.
 *
 * We use the data ready pin if it is available.  Otherwise, read the
 * status register.
 */
bool _data_ready(InertialSensor_Invensense_t *inven)
{
    if (inven->_drdy)
    {
        return (HAL_GPIO_ReadPin(inven->_drdy, inven->_drdy_Pin) == GPIO_PIN_SET);
    }
    uint8_t status = _register_read(inven, MPUREG_INT_STATUS);
    return (status & BIT_RAW_RDY_INT) != 0;
}

/*
 * Timer process to poll for new data from the Invensense. Called from bus thread with semaphore held
 */
void InertialSensor_Invensense_update(InertialSensor_Invensense_t *inven)
{
    if (inven->success)
    {
        if (inven->_accum.count && !osMutexWait(inven->mutex, 1))
        {
            float _fifo_accel_scale = inven->_accel_scale / inven->_accum.count;
            float _fifo_gyro_scale = GYRO_SCALE / inven->_accum.count;
            inven->accel[0] = inven->_accum.accel[0] * _fifo_accel_scale;
            inven->accel[1] = inven->_accum.accel[1] * _fifo_accel_scale;
            inven->accel[2] = inven->_accum.accel[2] * _fifo_accel_scale;
            inven->gyro[0] = inven->_accum.gyro[0] * _fifo_gyro_scale;
            inven->gyro[1] = inven->_accum.gyro[1] * _fifo_gyro_scale;
            inven->gyro[2] = inven->_accum.gyro[2] * _fifo_gyro_scale;

            inven->_accum.accel[0] = 0;
            inven->_accum.accel[1] = 0;
            inven->_accum.accel[2] = 0;
            inven->_accum.gyro[0] = 0;
            inven->_accum.gyro[1] = 0;
            inven->_accum.gyro[2] = 0;
            inven->_accum.count = 0;

            osMutexRelease(inven->mutex);
        }
    }
}

void InertialSensor_Invensense_loop(InertialSensor_Invensense_t *inven)
{
    while (true)
    {
        //imu->drdy->task = xTaskGetCurrentTaskHandle();
        //ulTaskNotifyTake(pdTRUE, 1);
        osDelay(1);
        if (!osMutexWait(inven->mutex, 0))
        {
            if (SPI_DEV_begin(inven->_dev, 0))
            {
                _read_fifo(inven);
                SPI_DEV_end(inven->_dev);
            }
            osMutexRelease(inven->mutex);
        }
    }
}

void InertialSensor_Invensense_monitor(InertialSensor_Invensense_t *inven)
{
    inven->pps = inven->cnt;
    inven->cnt = 0;
}

/*
  when doing fast sampling the sensor gives us 8k samples/second. Every 2nd accel sample is a duplicate.

  To filter this we first apply a 1p low pass filter at 188Hz, then we
  average over 8 samples to bring the data rate down to 1kHz. This
  gives very good aliasing rejection at frequencies well above what
  can be handled with 1kHz sample rates.
 */
bool _accumulate(InertialSensor_Invensense_t *inven, uint8_t *samples, uint8_t n_samples)
{
    int32_t tsum = 0;
    const int32_t clip_limit = AP_INERTIAL_SENSOR_ACCEL_CLIP_THRESH_MSS / inven->_accel_scale;
    bool clipped = false;
    bool ret = true;

    for (uint8_t i = 0; i < n_samples; i++)
    {
        const uint8_t *data = samples + MPU_SAMPLE_SIZE * i;

        // use temperatue to detect FIFO corruption
        int16_t t2 = int16_val(data, 3);
        if (!_check_raw_temp(inven, t2))
        {
            _fifo_reset(inven);
            //debug("[%s] temp reset %d %d", inven->name, inven->_raw_temp, t2);
            ret = false;
            break;
        }
        tsum += t2;

        float a[3];
        a[0] = int16_val(data, 1);
        a[1] = int16_val(data, 0);
        a[2] = -int16_val(data, 2);
        if (fabsf(a[0]) > clip_limit ||
            fabsf(a[0]) > clip_limit ||
            fabsf(a[0]) > clip_limit)
        {
            clipped = true;
        }
        inven->_accum.accel[0] += a[0];
        inven->_accum.accel[1] += a[1];
        inven->_accum.accel[2] += a[2];

        float g[3];
        g[0] = int16_val(data, 5);
        g[1] = int16_val(data, 4),
        g[2] = -int16_val(data, 6);

        inven->_accum.gyro[0] += g[0];
        inven->_accum.gyro[1] += g[1];
        inven->_accum.gyro[2] += g[2];
        inven->_accum.count++;
    }

    if (clipped)
    {
        ++inven->accel_clip_count;
    }

    inven->temp_degc = ((float)(tsum) / n_samples) * inven->temp_sensitivity + inven->temp_zero;

    return ret;
}

void _read_fifo(InertialSensor_Invensense_t *inven)
{
    uint8_t n_samples;
    uint16_t bytes_read;
    uint8_t *rx = inven->_fifo_buffer;
    bool need_reset = false;

    n_samples = 0;
    if (_block_read(inven, MPUREG_FIFO_COUNTH, rx, 2))
    {
        bytes_read = uint16_val(rx, 0);
        n_samples = bytes_read / MPU_SAMPLE_SIZE;
    }

    if (n_samples > 0)
    {
        /*
				testing has shown that if we have more than 32 samples in the
				FIFO then some of those samples will be corrupt. It always is
				the ones at the end of the FIFO, so clear those with a reset
				once we've read the first 24. Reading 24 gives us the normal
				number of samples for fast sampling at 400Hz

				On I2C with the much lower clock rates we need a lower threshold
				or we may never catch up
			 */
        if (n_samples > 32)
        {
            need_reset = true;
            n_samples = 24;
        }
        while (n_samples > 0)
        {
            uint8_t n = MIN(n_samples, MPU_FIFO_BUFFER_LEN);
            _block_read(inven, MPUREG_FIFO_R_W, rx, n * MPU_SAMPLE_SIZE);
            inven->cnt += n;
            inven->seq++;
            _accumulate(inven, rx, n);
            n_samples -= n;
        }

        if (need_reset)
        {
            //debug("[%s] fifo reset n_samples %u", inven->name, bytes_read / MPU_SAMPLE_SIZE);
            _fifo_reset(inven);
        }
    }
}

/*
  fetch temperature in order to detect FIFO sync errors
*/
bool _check_raw_temp(InertialSensor_Invensense_t *inven, int16_t t2)
{
    if (abs(t2 - inven->_raw_temp) < 400)
    {
        // cached copy OK
        return true;
    }
    uint8_t trx[2];
    if (_block_read(inven, MPUREG_TEMP_OUT_H, trx, 2))
    {
        inven->_raw_temp = int16_val(trx, 0);
    }
    return (abs(t2 - inven->_raw_temp) < 400);
}

bool _block_read(InertialSensor_Invensense_t *inven, uint8_t reg, uint8_t *buf,
                 uint32_t size)
{
    SPI_DEV_select(inven->_dev);
    bool rslt = SPI_DEV_read_registers(inven->_dev, reg, buf, size);
    SPI_DEV_unselect(inven->_dev);
    return rslt;
}

uint8_t _register_read(InertialSensor_Invensense_t *inven, uint8_t reg)
{
    uint8_t val = 0;
    SPI_DEV_select(inven->_dev);
    SPI_DEV_read_registers(inven->_dev, reg, &val, 1);
    SPI_DEV_unselect(inven->_dev);
    return val;
}

void _register_write(InertialSensor_Invensense_t *inven, uint8_t reg, uint8_t val)
{
    SPI_DEV_select(inven->_dev);
    SPI_DEV_write_register(inven->_dev, reg, val);
    SPI_DEV_unselect(inven->_dev);
}

/*
  set the DLPF filter frequency. Assumes caller has taken semaphore
 */
void _set_filter_register(InertialSensor_Invensense_t *inven)
{
    uint8_t config;

#if INVENSENSE_EXT_SYNC_ENABLE
    // add in EXT_SYNC bit if enabled
    config = (MPUREG_CONFIG_EXT_SYNC_AZ << MPUREG_CONFIG_EXT_SYNC_SHIFT);
#else
    config = 0;
#endif

    /* set divider for internal sample rate to 0x1F when fast
             sampling enabled. This reduces the impact of the slave
             sensor on the sample rate. It ends up with around 75Hz
             slave rate, and reduces the impact on the gyro and accel
             sample rate, ending up with around 7760Hz gyro rate and
             3880Hz accel rate
             */
    _register_write(inven, MPUREG_I2C_SLV4_CTRL, 0x1F);

    // this gives us 8kHz sampling on gyros and 4kHz on accels
    config |= BITS_DLPF_CFG_256HZ_NOLPF2;

    config |= MPUREG_CONFIG_FIFO_MODE_STOP;
    _register_write(inven, MPUREG_CONFIG, config);

    if (inven->_mpu_type != Invensense_MPU6000)
    {
        // setup for 4kHz accels
        _register_write(inven, ICMREG_ACCEL_CONFIG2, ICM_ACC_FCHOICE_B);
    }
}

/*
  check whoami for sensor type
 */
bool _check_whoami(InertialSensor_Invensense_t *inven)
{
    uint8_t whoami = _register_read(inven, MPUREG_WHOAMI);
    switch (whoami)
    {
    case MPU_WHOAMI_6000:
        inven->_mpu_type = Invensense_MPU6000;
        return true;
    case MPU_WHOAMI_6500:
        inven->_mpu_type = Invensense_MPU6500;
        return true;
    case MPU_WHOAMI_MPU9250:
    case MPU_WHOAMI_MPU9255:
        inven->_mpu_type = Invensense_MPU9250;
        return true;
    case MPU_WHOAMI_20608:
        inven->_mpu_type = Invensense_ICM20608;
        return true;
    case MPU_WHOAMI_20600:
        info("Find Invensense_ICM20600");
        inven->_mpu_type = Invensense_ICM20600;
        return true;
    case MPU_WHOAMI_20602:
        info("Find Invensense_ICM20602");
        inven->_mpu_type = Invensense_ICM20602;
        return true;
    case MPU_WHOAMI_20601:
        info("Find Invensense_ICM20601");
        inven->_mpu_type = Invensense_ICM20601;
        return true;
    case MPU_WHOAMI_ICM20789:
    case MPU_WHOAMI_ICM20789_R1:
        inven->_mpu_type = Invensense_ICM20789;
        return true;
    case MPU_WHOAMI_ICM20689:
        info("Find Invensense_ICM20689");
        inven->_mpu_type = Invensense_ICM20689;
        return true;
    }
    // not a value WHOAMI result
    info("not a value WHOAMI result");
    return false;
}

bool _hardware_init(InertialSensor_Invensense_t *inven)
{
    if (!SPI_DEV_begin(inven->_dev, 3000))
    {
        return false;
    }

    if (!_check_whoami(inven))
    {
        goto fail;
    }

    // Chip reset
    uint8_t tries;
    for (tries = 0; tries < 5; tries++)
    {
        inven->_last_stat_user_ctrl = _register_read(inven, MPUREG_USER_CTRL);

        /* First disable the master I2C to avoid hanging the slaves on the
         * aulixiliar I2C bus - it will be enabled again if the AuxiliaryBus
         * is used */
        if (inven->_last_stat_user_ctrl & BIT_USER_CTRL_I2C_MST_EN)
        {
            inven->_last_stat_user_ctrl &= ~BIT_USER_CTRL_I2C_MST_EN;
            _register_write(inven, MPUREG_USER_CTRL, inven->_last_stat_user_ctrl);
            vTaskDelay(10);
        }

        /* reset device */
        _register_write(inven, MPUREG_PWR_MGMT_1, BIT_PWR_MGMT_1_DEVICE_RESET);
        vTaskDelay(100);

        /* bus-dependent initialization */
        if (true)
        {
            /* Disable I2C bus if SPI selected (Recommended in Datasheet to be
             * done just after the device is reset) */
            inven->_last_stat_user_ctrl |= BIT_USER_CTRL_I2C_IF_DIS;
            _register_write(inven, MPUREG_USER_CTRL, inven->_last_stat_user_ctrl);
        }

        // Wake up device and select GyroZ clock. Note that the
        // Invensense starts up in sleep mode, and it can take some time
        // for it to come out of sleep
        _register_write(inven, MPUREG_PWR_MGMT_1, BIT_PWR_MGMT_1_CLK_ZGYRO);
        vTaskDelay(5);

        // check it has woken up
        if (_register_read(inven, MPUREG_PWR_MGMT_1) == BIT_PWR_MGMT_1_CLK_ZGYRO)
        {
            info("[%s] has woken up.", inven->name);
            break;
        }

        vTaskDelay(10);
        if (_data_ready(inven))
        {
            info("[%s] data_ready.", inven->name);
            break;
        }
    }

    if (tries == 5)
    {
        info("[%s] Failed to boot Invensense 5 times", inven->name);
        goto fail;
    }

    if (inven->_mpu_type == Invensense_ICM20608 ||
				inven->_mpu_type == Invensense_ICM20602 ||
        inven->_mpu_type == Invensense_ICM20600)
    {
        // this avoids a sensor bug, see description above
        _register_write(inven, MPUREG_ICM_UNDOC1, MPUREG_ICM_UNDOC1_VALUE);
    }

    if (_data_ready(inven))
    {
        info("[%s] data_ready.", inven->name);
    }
    SPI_DEV_end(inven->_dev);
    return true;
fail:
    SPI_DEV_end(inven->_dev);
    return false;
}