// lis3dsh数据手册：https://www.findic.com/doc/browser/8LWZ55XQZ?doc_id=48813340#locale=zh-CN
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/i2c.h"
#include "driver/gpio.h"
#include "esp_dsp.h"
#include "esp_err.h"
#include "esp_log.h"

#include "shake_detect.h"
#include "lis3dsh_lib/lis3dsh_reg.h"

static const char *TAG = "shake_detect";

static QueueHandle_t int1_evt_queue = NULL;

// 0--SPI, 1--I2C
static void lis3dsh_cs_io_init(int mode)
{
    gpio_config_t io_conf = {};
    io_conf.intr_type = GPIO_INTR_DISABLE;
    io_conf.mode = GPIO_MODE_OUTPUT;
    io_conf.pin_bit_mask = (1ULL << LIS3DSH_CS_IO);
    io_conf.pull_down_en = 0;
    io_conf.pull_up_en = 0;
    gpio_config(&io_conf);
    gpio_set_level(LIS3DSH_CS_IO, mode);
}

static void lis3dsh_i2c_sel_io_init(int addr01)
{
    gpio_config_t io_conf = {};
    io_conf.intr_type = GPIO_INTR_DISABLE;
    io_conf.mode = GPIO_MODE_OUTPUT;
    io_conf.pin_bit_mask = (1ULL << LIS3DSH_I2C_SEL_IO);
    io_conf.pull_down_en = 0;
    io_conf.pull_up_en = 0;
    gpio_config(&io_conf);
    gpio_set_level(LIS3DSH_I2C_SEL_IO, addr01);
}

static esp_err_t lis3dsh_i2c_master_init(void)
{
    int i2c_master_port = LIS3DSH_I2C_MASTER_NUM;

    lis3dsh_cs_io_init(1);
    lis3dsh_i2c_sel_io_init(0);

    i2c_config_t conf = {
        .mode = I2C_MODE_MASTER,
        .sda_io_num = LIS3DSH_I2C_MASTER_SDA_IO,
        .scl_io_num = LIS3DSH_I2C_MASTER_SCL_IO,
        .sda_pullup_en = GPIO_PULLUP_ENABLE,
        .scl_pullup_en = GPIO_PULLUP_ENABLE,
        .master.clk_speed = LIS3DSH_I2C_MASTER_FREQ_HZ,
    };

    i2c_param_config(i2c_master_port, &conf);

    return i2c_driver_install(i2c_master_port, conf.mode, 0, 0, 0);
}

static void lis3dsh_led_init(void)
{
    gpio_config_t io_conf = {};
    io_conf.intr_type = GPIO_INTR_DISABLE;
    io_conf.mode = GPIO_MODE_OUTPUT;
    io_conf.pin_bit_mask = (1ULL << LIS3DSH_LED_IO);
    io_conf.pull_down_en = 0;
    io_conf.pull_up_en = 0;
    gpio_config(&io_conf);
}

static void lis3dsh_led_toggle(void)
{
    static int state = 0;
    gpio_set_level(LIS3DSH_LED_IO, state);
    state = !state;
}

static void IRAM_ATTR int1_isr_handler(void* arg)
{
    uint32_t gpio_num = (uint32_t) arg;
    if (gpio_num == LIS3DSH_INT1_IO)
    {
        xQueueSendFromISR(int1_evt_queue, &gpio_num, NULL);
    }
}

static void lis3dsh_int1_int_init(void)
{
    gpio_config_t io_conf = {};
    io_conf.intr_type = GPIO_INTR_NEGEDGE;  
    io_conf.mode = GPIO_MODE_INPUT;
    io_conf.pin_bit_mask = (1ULL << LIS3DSH_INT1_IO);
    io_conf.pull_down_en = 0;
    io_conf.pull_up_en = 1;
    gpio_config(&io_conf);

    int1_evt_queue = xQueueCreate(1, sizeof(uint32_t));

    gpio_install_isr_service(0);   
    gpio_isr_handler_add(LIS3DSH_INT1_IO, int1_isr_handler, (void*) LIS3DSH_INT1_IO);
}

static int platform_init(void)
{
    ESP_ERROR_CHECK(lis3dsh_i2c_master_init());
    lis3dsh_int1_int_init();
    lis3dsh_led_init();
    return LIS3DSH_I2C_MASTER_NUM;
}

static int32_t platform_write(void *handle, uint8_t reg, uint8_t *bufp, uint16_t len)
{
    uint8_t write_buf[128];
    if (len > 127) 
        return -1;
    write_buf[0] = reg;
    memcpy(&write_buf[1], bufp, len);
    return i2c_master_write_to_device(LIS3DSH_I2C_MASTER_NUM, LIS3DSH_I2C_ADD_L >> 1, write_buf, len + 1, LIS3DSH_I2C_MASTER_TIMEOUT_MS / portTICK_PERIOD_MS);
}

static int32_t platform_read(void *handle, uint8_t reg, uint8_t *bufp, uint16_t len)
{
    return i2c_master_write_read_device(LIS3DSH_I2C_MASTER_NUM, LIS3DSH_I2C_ADD_L >> 1, &reg, 1, bufp, len, LIS3DSH_I2C_MASTER_TIMEOUT_MS / portTICK_PERIOD_MS);
}

static void platform_delay(uint32_t ms)
{
    vTaskDelay(pdMS_TO_TICKS(ms));
}

stmdev_ctx_t dev_ctx;
void lis3dsh_init(void)
{
    int ret = 0;
    lis3dsh_pin_int1_route_t int1_route;
    lis3dsh_all_sources_t all_sources;
    lis3dsh_bus_mode_t bus_mode;
    lis3dsh_status_var_t status;
    lis3dsh_ctrl_reg3_t ctrl_reg3;
    lis3dsh_id_t id = {0};
    lis3dsh_md_t md;
    lis3dsh_reg_t reg;

    /* Initialize mems driver interface */
    dev_ctx.handle = (void *)platform_init();
    dev_ctx.write_reg = platform_write;
    dev_ctx.read_reg = platform_read;
    
    /* Wait sensor boot time */
    platform_delay(BOOT_TIME);
    /* Check device ID */
    ret = lis3dsh_id_get(&dev_ctx, &id);
    ESP_LOGI(TAG, "ret=%d, lis3dsh_id=%02x%02x%02x", ret, id.whoami, id.info1, id.info2);
    while (id.whoami != LIS3DSH_ID)
    {
        platform_delay(100);
        memset(&id, 0, sizeof(id));
        ret = lis3dsh_id_get(&dev_ctx, &id);
        ESP_LOGI(TAG, "ret=%d, lis3dsh_id=%02x%02x%02x", ret, id.whoami, id.info1, id.info2);
    }

    // 有时候系统启动会卡在这个循环里
    // /* Restore default configuration */
    // lis3dsh_init_set(&dev_ctx, LIS3DSH_RESET);
    // do
    // {
    //     lis3dsh_status_get(&dev_ctx, &status);
    // } while (status.sw_reset);

    /* Set bdu and if_inc recommended for driver usage */
    lis3dsh_init_set(&dev_ctx, LIS3DSH_DRV_RDY);

    /* Set Output Data Rate */
    // ctrl_reg4, ctrl_reg5
    lis3dsh_mode_get(&dev_ctx, &md);
    md.fs = LIS3DSH_2g;
    md.odr = LIS3DSH_1kHz6;
    lis3dsh_mode_set(&dev_ctx, &md);

    /* FIFO configuration */
    // fifo_ctrl
    lis3dsh_read_reg(&dev_ctx, LIS3DSH_FIFO_CTRL, &reg.byte, 1);
    reg.fifo_ctrl.fmode = 2; /* FIFO stream mode */
    reg.fifo_ctrl.wtmp = 16;
    lis3dsh_write_reg(&dev_ctx, LIS3DSH_FIFO_CTRL, &reg.byte, 1);

    // ctrl_reg6:fifo使能，地址自增使能，fifo水位中断int1使能
    lis3dsh_read_reg(&dev_ctx, LIS3DSH_CTRL_REG6, &reg.byte, 1);
    memset(&reg.byte, 0, sizeof(reg.byte));
    ESP_LOGI(TAG, "LIS3DSH_CTRL_REG6 old=%02x", reg.byte);
    reg.ctrl_reg6.fifo_en = PROPERTY_ENABLE;
    reg.ctrl_reg6.wtm_en = PROPERTY_ENABLE;
    reg.ctrl_reg6.add_inc = PROPERTY_ENABLE;
    reg.ctrl_reg6.p1_wtm = PROPERTY_ENABLE;
    lis3dsh_write_reg(&dev_ctx, LIS3DSH_CTRL_REG6, &reg.byte, 1);
    lis3dsh_read_reg(&dev_ctx, LIS3DSH_CTRL_REG6, &reg.byte, 1);
    ESP_LOGI(TAG, "LIS3DSH_CTRL_REG6 new=%02x", reg.byte);

    // ctrl_reg3:中断配置：int1使能，DRDY失能，int2失能，低电平有效，锁存 
    lis3dsh_read_reg(&dev_ctx, LIS3DSH_CTRL_REG3, (uint8_t *)&ctrl_reg3, 1);
    memset(&ctrl_reg3, 0, sizeof(ctrl_reg3));
    ctrl_reg3.int1_en = PROPERTY_ENABLE;
    lis3dsh_write_reg(&dev_ctx, LIS3DSH_CTRL_REG3, (uint8_t *)&ctrl_reg3, 1);

    lis3dsh_read_reg(&dev_ctx, LIS3DSH_CTRL_REG4, &reg.byte, 1);
    ESP_LOGI(TAG, "LIS3DSH_CTRL_REG4 old=%02x", reg.byte);
}



// 乒乓缓冲区定义，fft一次需要取N_SAMPLES个数据运算，定义缓冲区大小为其2倍，在写另一边的时候fft可以对这边数据进行运算
__attribute__((aligned(16)))
lis3dsh_a_data_t lis3dsh_data[N_SAMPLES * 2];
int buf_r = 0;
int fft_h = 0;
int angle_h = 0;
static QueueHandle_t fft_queue = NULL;
static QueueHandle_t angle_queue = NULL;

void lis3dsh_read_task(void)
{
    uint32_t io_num;
    lis3dsh_reg_t reg;
    uint8_t rx_buffer[128];

    while (1)
    {
        if(xQueueReceive(int1_evt_queue, &io_num, pdMS_TO_TICKS(1000))) 
        {
            lis3dsh_read_reg(&dev_ctx, LIS3DSH_FIFO_SRC, &reg.byte, 1);
            // ESP_LOGI(TAG, "fss:%d", reg.fifo_src.fss);

            for (int i = 0; i < reg.fifo_src.fss; i++)
            {
                lis3dsh_read_reg(&dev_ctx, LIS3DSH_OUT_X_L, rx_buffer, 6);
                // ESP_LOGI(TAG, "read %02d:%10f, %10f, %10f", i, (int16_t)(rx_buffer[0] | (rx_buffer[1] << 8)) * 2.0 / 0x7fff, (int16_t)(rx_buffer[2] | (rx_buffer[3] << 8)) * 2.0 / 0x7fff, (int16_t)(rx_buffer[4] | (rx_buffer[5] << 8)) * 2.0 / 0x7fff);
            
                lis3dsh_data[buf_r].a[0] = (int16_t)(rx_buffer[0] | (rx_buffer[1] << 8)) * 2.0 / 0x7fff;
                lis3dsh_data[buf_r].a[1] = (int16_t)(rx_buffer[2] | (rx_buffer[3] << 8)) * 2.0 / 0x7fff;   // y 有重力加速度分量，计算fft的时候需要去掉
                lis3dsh_data[buf_r].a[2] = (int16_t)(rx_buffer[4] | (rx_buffer[5] << 8)) * 2.0 / 0x7fff;
                buf_r++;

                if (buf_r == N_SAMPLES || buf_r == N_SAMPLES * 2)
                {
                    // ESP_LOGI(TAG, "send fft queue");
                    fft_h = buf_r - N_SAMPLES;
                    xQueueSend(fft_queue, &fft_h, portMAX_DELAY);
                }
                if (buf_r != 0 && buf_r % ANGLE_N_SAMPLES == 0)
                {
                    // ESP_LOGI(TAG, "send angle queue");
                    angle_h = buf_r - ANGLE_N_SAMPLES;
                    xQueueSend(angle_queue, &angle_h, portMAX_DELAY);
                }
                if (buf_r == N_SAMPLES * 2)
                {
                    buf_r = 0;
                }
            }
        }
        else 
        {
            ESP_LOGI(TAG, "wait int1...");  // 很奇怪，第一次中断进不去，只有读一次才能进
            lis3dsh_read_reg(&dev_ctx, LIS3DSH_OUT_X_L, rx_buffer, 6);
            ESP_LOGI(TAG, "read :%f, %f, %f", (int16_t)(rx_buffer[0] | (rx_buffer[1] << 8)) * 2.0 / 0x7fff, (int16_t)(rx_buffer[2] | (rx_buffer[3] << 8)) * 2.0 / 0x7fff, (int16_t)(rx_buffer[4] | (rx_buffer[5] << 8)) * 2.0 / 0x7fff);
        }
    }
}


shake_detect_cfg_t shake_cfg = {10, 25, 0.06, 0.03, 2, 3, 3};
__attribute__((aligned(16)))
fft_data_t fft_data;

bool shake_detect_fun(float *feq_data)
{
    if (feq_data == NULL) return false;

    float a_max = 0;
    float a_th = fft_data.is_shake ? shake_cfg.exit_amplitude_th : shake_cfg.enter_amplitude_th;
    for (int i = shake_cfg.feq_min_id; i <= shake_cfg.feq_max_id; i++)
    {
        a_max = a_max >= feq_data[i] ? a_max : feq_data[i];
    }
    if (a_max >= a_th)
        return true;
    return false;
}

void fft_calculate_task(void)
{
    esp_err_t ret;
    fft_data.y1_cf = &fft_data.y_cf[0];
    fft_data.y2_cf = &fft_data.y_cf[N_SAMPLES];

    // fft 初始化
    ret = dsps_fft2r_init_fc32(NULL, CONFIG_DSP_MAX_FFT_SIZE);
    if (ret  != ESP_OK) {
        ESP_LOGI(TAG, "Not possible to initialize FFT. Error = %i", ret);
        return;
    }
    // Generate hann window
    dsps_wind_hann_f32(fft_data.wind, N_SAMPLES);
    
    while (1)
    {
        // ESP_LOGI(TAG, "fft wait...");
        if(xQueueReceive(fft_queue, &fft_h, portMAX_DELAY)) 
        {
            // Convert two input vectors to one complex vector
            uint8_t ch1 = shake_cfg.channel1;
            uint8_t ch2 = shake_cfg.channel2;
            if (ch1 <= 0 || ch1 > 3 || ch2 <= 0 || ch2 > 3)
            {
                ESP_LOGI(TAG, "CH select err[1-3]: ch1=%d, ch2=%d", ch1, ch2);
                continue;
            }
            for (int i = 0 ; i < N_SAMPLES ; i++) {
                fft_data.y_cf[i * 2 + 0] = lis3dsh_data[fft_h + i].a[ch1 - 1] * fft_data.wind[i];
                fft_data.y_cf[i * 2 + 1] = lis3dsh_data[fft_h + i].a[ch2 - 1] * fft_data.wind[i];
            }
            // FFT
            unsigned int start_b = dsp_get_cpu_cycle_count();
            dsps_fft2r_fc32(fft_data.y_cf, N_SAMPLES);
            unsigned int end_b = dsp_get_cpu_cycle_count();
            // Bit reverse
            dsps_bit_rev_fc32(fft_data.y_cf, N_SAMPLES);
            // Convert one complex vector to two complex vectors
            dsps_cplx2reC_fc32(fft_data.y_cf, N_SAMPLES);

            // 归一化处理参考：https://blog.csdn.net/weixin_39591031/article/details/110392352
            for (int i = 0 ; i < N_SAMPLES / 2 ; i++) {
                int normal_div = N_SAMPLES / 2.0;
                if (i == 0) normal_div = N_SAMPLES;
                fft_data.y1_cf[i] = (sqrt(fft_data.y1_cf[i * 2 + 0] * fft_data.y1_cf[i * 2 + 0] + fft_data.y1_cf[i * 2 + 1] * fft_data.y1_cf[i * 2 + 1]) / normal_div);
                fft_data.y2_cf[i] = (sqrt(fft_data.y2_cf[i * 2 + 0] * fft_data.y2_cf[i * 2 + 0] + fft_data.y2_cf[i * 2 + 1] * fft_data.y2_cf[i * 2 + 1]) / normal_div);
                // Simple way to show two power spectrums as one plot
                fft_data.sum_y[i] = fmax(fft_data.y1_cf[i], fft_data.y2_cf[i]);
                float freq = i * 1.0 / N_SAMPLES * 1600;  // 单位HZ
                // ESP_LOGI(TAG, "fft %d: f=%f, y=%f, z=%f sum=%f", i, freq, fft_data.y1_cf[i], fft_data.y2_cf[i], fft_data.sum_y[i]);
            }

            ESP_LOGW(TAG, "Signal ch1");
            dsps_view(fft_data.y1_cf, N_SAMPLES / 2, 64, 10,  -2, 2, '|');
            ESP_LOGI(TAG, "FFT for %i complex points take %i cycles", N_SAMPLES, end_b - start_b);

            // 震动检测
            float *feq_data = NULL;
            switch (shake_cfg.select)
            {
                case 1:
                    feq_data = fft_data.y1_cf;
                    break;
                case 2:
                    feq_data = fft_data.y2_cf;
                    break;
                case 3:
                    feq_data = fft_data.sum_y;
                    break;
                default:
                    ESP_LOGI(TAG, "select err[1-3]: %d", shake_cfg.select);
                    continue;
                    break;
            }
            fft_data.is_shake = shake_detect_fun(feq_data);
            // ESP_LOGI(TAG, "is_shake = %d", fft_data.is_shake);
            lis3dsh_led_toggle();
        }
    }
}

#define PI (3.1415926)
void lis3dsh_calculating_angle(float ax, float ay, float az, float *anglex, float *angley, float *anglez)
{
    float axy = sqrt(ax * ax + ay * ay);
    *anglez = (atan2(axy, az) * 180 / PI);

    float axz = sqrt(ax * ax + az * az);
    *angley = (atan2(ay, axz) * 180 / PI);

    float ayz = sqrt(ay * ay + az * az);
    *anglex = (atan2(ax, ayz) * 180 / PI);
    
}

angle_data_t angle_data;
void angle_calculate_task(void)
{
    float ax = 0, ay = 0, az = 0;

    while (1)
    {
        // ESP_LOGI(TAG, "angle wait...");
        if(xQueueReceive(angle_queue, &angle_h, portMAX_DELAY)) 
        {
            for (int i = 0; i < ANGLE_N_SAMPLES; i++)
            {
                ax += lis3dsh_data[angle_h + i].a[0];
                ay += lis3dsh_data[angle_h + i].a[1];
                az += lis3dsh_data[angle_h + i].a[2];
            }
            ax /= ANGLE_N_SAMPLES;
            ay /= ANGLE_N_SAMPLES;
            az /= ANGLE_N_SAMPLES;
            lis3dsh_calculating_angle(ax, ay, az, &angle_data.x, &angle_data.y, &angle_data.z);

            // ESP_LOGI(TAG, "angle: x=%5.5f\ty=%5.5f\tz=%5.5f", angle_data.x, angle_data.y, angle_data.z);
        }
          
    }

}


void shake_detect_init(void)
{
    lis3dsh_init();

    fft_queue = xQueueCreate(1, sizeof(uint32_t));
    angle_queue = xQueueCreate(1, sizeof(uint32_t));


    // 震动检测配置参数加载

    xTaskCreatePinnedToCore(lis3dsh_read_task, "lis3dsh_read_task", 4096, NULL, 9, NULL, tskNO_AFFINITY);
    xTaskCreatePinnedToCore(fft_calculate_task, "fft_calculate_task", 4096, NULL, 9, NULL, tskNO_AFFINITY);
    xTaskCreatePinnedToCore(angle_calculate_task, "angle_calculate_task", 4096, NULL, 9, NULL, tskNO_AFFINITY);
}