/*
	Copyright 2020 Mitch Lustig
	Copyright 2026 Benjamin Vedder benjamin@vedder.se

	This file is part of the VESC firmware.

	The VESC firmware is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	The VESC firmware is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
	*/

#include "lsm6ds3.h"
#include "terminal.h"
#include "commands.h"
#include "utils_math.h"

#include <stdio.h>
#include <string.h>

static thread_t *lsm6ds3_thread_ref = NULL;
static i2c_bb_state *m_i2c_bb;
static spi_bb_state *m_spi_bb;
static SPIDriver *m_hwspi_dev;
static volatile uint16_t lsm6ds3_addr;
static int rate_hz = 1000;
static IMU_FILTER filter;
static const uint8_t spi_tx_0_12[12] = {0};

static void terminal_read_reg(int argc, const char **argv);
static bool read_reg(uint8_t reg, uint8_t *res);
static bool write_reg(uint8_t reg, uint8_t value);
static bool read_gyro_accel(uint8_t *res);
static THD_FUNCTION(lsm6ds3_thread, arg);

// Function pointers
static void(*read_callback)(float *accel, float *gyro, float *mag) = 0;

void lsm6ds3_set_rate_hz(int hz) {
	rate_hz = hz;
}

void lsm6ds3_set_filter(IMU_FILTER f) {
	filter = f;
}

void lsm6ds3_init(i2c_bb_state *i2c_state, spi_bb_state *spi_state, SPIDriver *spi_hw,
		stkalign_t *work_area, size_t work_area_size) {

	read_callback = 0;

	m_i2c_bb = i2c_state;
	m_spi_bb = spi_state;
	m_hwspi_dev = spi_hw;

	uint8_t rxb[1];

	lsm6ds3_addr = LSM6DS3_ACC_GYRO_ADDR_A;
	bool res = read_reg(LSM6DS3_ACC_GYRO_WHO_AM_I_REG, rxb);
	if (!res || (rxb[0] != 0x69 && rxb[0] != 0x6A && rxb[0] != 0x6C)) {
		commands_printf("LSM6DS3 Address A failed, trying B (rx: %d)", rxb[0]);
		lsm6ds3_addr = LSM6DS3_ACC_GYRO_ADDR_B;
		res = read_reg(LSM6DS3_ACC_GYRO_WHO_AM_I_REG, rxb);
		if (!res || (rxb[0] != 0x69 && rxb[0] != 0x6A && rxb[0] != 0x6C)) {
			commands_printf("LSM6DS3 Address B failed (rx: %d)", rxb[0]);
			return;
		}
	}

	bool is_trc = false;
	if (rxb[0] == 0x6A){
		is_trc = true;
	}

	// Accelerometer resolution and data rate
	uint8_t regv = LSM6DS3_ACC_GYRO_FS_XL_16g;
	regv |= LSM6DS3_ACC_GYRO_ODR_XL_6660Hz;

	// Accelerometer filtering
	#define LSM6DS3TRC_BW0_XL 0x1
	#define LSM6DS3TRC_LPF1_BW_SEL 0x2
	if (is_trc) {
		// Always use accelerometer analog low-pass at 400Hz
		regv |= LSM6DS3TRC_BW0_XL;
	} else if (rate_hz >= 208 && filter >= IMU_FILTER_MEDIUM) {
		// Filter at ODR/4 for MEDIUM and ODR/8 for HIGH
		// This filter also needs to be enabled in CTRL4_C
		int scaled_rate = filter == IMU_FILTER_HIGH ? rate_hz / 2 : rate_hz;
		if (scaled_rate <= 208) {
			regv |= LSM6DS3_ACC_GYRO_BW_XL_50Hz;
		} else if (scaled_rate <= 416) {
			regv |= LSM6DS3_ACC_GYRO_BW_XL_100Hz;
		} else if (scaled_rate <= 833) {
			regv |= LSM6DS3_ACC_GYRO_BW_XL_200Hz;
		}
	}

	res = write_reg(LSM6DS3_ACC_GYRO_CTRL1_XL, regv);
	if (!res){
		commands_printf("LSM6DS3 Accel Config FAILED");
		return;
	}

	// Extra accelerometer filtering for TRC variant
	if (is_trc) {
		#define LSM6DS3TRC_LPF2_XL_EN 0x80
		#define LSM6DS3TRC_HPCF_XL_ODR9 0x40
		#define LSM6DS3TRC_HPCF_XL_ODR50 0x00
		#define LSM6DS3TRC_HPCF_XL_ODR100 0x20
		regv = 0;
		if (filter == IMU_FILTER_MEDIUM) {
			regv |= LSM6DS3TRC_LPF2_XL_EN | LSM6DS3TRC_HPCF_XL_ODR50;
		} else if (filter == IMU_FILTER_HIGH) {
			regv |= LSM6DS3TRC_LPF2_XL_EN | LSM6DS3TRC_HPCF_XL_ODR100;
		}

		res = write_reg(LSM6DS3_ACC_GYRO_CTRL8_XL, regv);
		if (!res) {
			commands_printf("LSM6DS3 Accel Filter Config FAILED");
			return;
		}
	}

	// Gyro resolution and data rate
	regv = LSM6DS3_ACC_GYRO_FS_G_2000dps;

	if (is_trc) {
		regv |= LSM6DS3TRC_ACC_GYRO_ODR_G_6660Hz;
	} else {
		// On non-TRC there is no dedicated configurable gyro filter, the filtering
		// seems to depend on the actual ODR, so we can't oversample it.
		if (rate_hz <= 13) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_13Hz;
		} else if (rate_hz <= 26) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_26Hz;
		} else if (rate_hz <= 52) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_52Hz;
		} else if (rate_hz <= 104) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_104Hz;
		} else if (rate_hz <= 208) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_208Hz;
		} else if (rate_hz <= 416) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_416Hz;
		} else if (rate_hz <= 833) {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_833Hz;
		} else {
			regv |= LSM6DS3_ACC_GYRO_ODR_G_1660Hz;
		}
	}
	res = write_reg(LSM6DS3_ACC_GYRO_CTRL2_G, regv);
	if (!res){
		commands_printf("LSM6DS3 Gyro Config FAILED");
		return;
	}

	// Extra gyro filtering for TRC variant
	if (is_trc) {
		#define LSM6DS3TRC_FTYPE_L 0x00
		#define LSM6DS3TRC_FTYPE_M 0x01
		#define LSM6DS3TRC_FTYPE_H 0x10
		regv = 0;
		if (filter == IMU_FILTER_LOW) {
			regv |= LSM6DS3TRC_FTYPE_L;
		} else if (filter == IMU_FILTER_MEDIUM) {
			regv |= LSM6DS3TRC_FTYPE_M;
		} else if (filter == IMU_FILTER_HIGH) {
			regv |= LSM6DS3TRC_FTYPE_H;
		}

		res = write_reg(LSM6DS3_ACC_GYRO_CTRL6_G, regv);
		if (!res){
			commands_printf("LSM6DS3 Gyro Filter FAILED");
			return;
		}
	}

	// Miscellaneous filtering configuration in CTRL4_C
	// TRC Variant CTRL4 register is very different from other variants
	regv = 0;
	if (is_trc) {
		// Enable gyroscope digital low-pass filter LPF1
		regv = LSM6DS3_ACC_GYRO_LPF1_SEL_G_ENABLED;
	} else if (rate_hz >= 208 && filter >= IMU_FILTER_MEDIUM) {
		// Standard LSM6DS3 only: Set XL anti-aliasing filter to be manually configured
		regv = LSM6DS3_ACC_GYRO_BW_SCAL_ODR_ENABLED;
	}
	res = write_reg(LSM6DS3_ACC_GYRO_CTRL4_C, regv);
	if (!res) {
		commands_printf("LSM6DS3 Misc Filter Config FAILED");
		return;
	}

	// Configure block update and register auto-increment
	regv = LSM6DS3_ACC_GYRO_BDU_BLOCK_UPDATE | LSM6DS3_ACC_GYRO_IF_INC_ENABLED;
	res = write_reg(LSM6DS3_ACC_GYRO_CTRL3_C, regv);
	if (!res) {
		commands_printf("LSM6DS3 BDU Config FAILED");
		return;
	}

	terminal_register_command_callback(
			"lsm_read_reg",
			"Read register of the LSM6DS3",
			"[reg]",
			terminal_read_reg);

	lsm6ds3_thread_ref = chThdCreateStatic(work_area, work_area_size, NORMALPRIO, lsm6ds3_thread, NULL);
}

void lsm6ds3_stop(void) {
	if (lsm6ds3_thread_ref != NULL){
		chThdTerminate(lsm6ds3_thread_ref);
		chThdWait(lsm6ds3_thread_ref);
	}
	lsm6ds3_thread_ref = NULL;
	terminal_unregister_callback(terminal_read_reg);
}

void lsm6ds3_set_read_callback(void(*func)(float *accel, float *gyro, float *mag)) {
	read_callback = func;
}

static bool read_reg(uint8_t reg, uint8_t *res) {
	bool ok = false;

	if (m_i2c_bb) {
		uint8_t txb[1];
		uint8_t rxb[1];
		txb[0] = reg;
		ok = i2c_bb_tx_rx(m_i2c_bb, lsm6ds3_addr, txb, 1, rxb, 1);

		if (ok) {
			*res = rxb[0];
		} else {
			*res = 0;
		}
	} else if (m_spi_bb) {
		chMtxLock(&(m_spi_bb->mutex));
		spi_bb_begin(m_spi_bb);
		spi_bb_exchange_8_mode_3(m_spi_bb, reg | 0x80);
		spi_bb_delay();
		*res = spi_bb_exchange_8_mode_3(m_spi_bb, 0);
		spi_bb_end(m_spi_bb);
		chMtxUnlock(&(m_spi_bb->mutex));
		ok = true;
	} else if (m_hwspi_dev) {
		// Use polled exchange here as the DMA-version seems to freeze
		// the CAN process thread for some reason. Performance does not
		// matter during config.
		spiAcquireBus(m_hwspi_dev);
		spiSelect(m_hwspi_dev);
		spiPolledExchange(m_hwspi_dev, reg | 0x80);
		spi_bb_delay();
		*res = spiPolledExchange(m_hwspi_dev, 0);
		spiUnselect(m_hwspi_dev);
		spiReleaseBus(m_hwspi_dev);
		ok = true;
	}

	return ok;
}

static bool write_reg(uint8_t reg, uint8_t value) {
	bool ok = false;

	if (m_i2c_bb) {
		uint8_t txb[2];
		uint8_t rxb[1];
		txb[0] = reg;
		txb[1] = value;
		ok = i2c_bb_tx_rx(m_i2c_bb, lsm6ds3_addr, txb, 2, rxb, 1);
	} else if (m_spi_bb) {
		chMtxLock(&(m_spi_bb->mutex));
		spi_bb_begin(m_spi_bb);
		spi_bb_exchange_8_mode_3(m_spi_bb, reg & 0x7F);
		spi_bb_delay();
		spi_bb_exchange_8_mode_3(m_spi_bb, value);
		spi_bb_end(m_spi_bb);
		chMtxUnlock(&(m_spi_bb->mutex));
		ok = true;
	} else if (m_hwspi_dev) {
		// Use polled exchange here as the DMA-version seems to freeze
		// the CAN process thread for some reason. Performance does not
		// matter during config.
		spiAcquireBus(m_hwspi_dev);
		spiSelect(m_hwspi_dev);
		spiPolledExchange(m_hwspi_dev, reg & 0x7f);
		spi_bb_delay();
		spiPolledExchange(m_hwspi_dev, value);
		spiUnselect(m_hwspi_dev);
		spiReleaseBus(m_hwspi_dev);
		ok = true;
	}

	return ok;
}

static bool read_gyro_accel(uint8_t *res) {
	bool ok = false;

	if (m_i2c_bb) {
		uint8_t txb[1];
		txb[0] = LSM6DS3_ACC_GYRO_OUTX_L_G;
		ok = i2c_bb_tx_rx(m_i2c_bb, lsm6ds3_addr, txb, 1, res, 12);

		if (!ok) {
			memset(res, 0, 12);
		}
	} else if (m_spi_bb) {
		chMtxLock(&(m_spi_bb->mutex));
		spi_bb_begin(m_spi_bb);
		spi_bb_exchange_8_mode_3(m_spi_bb, LSM6DS3_ACC_GYRO_OUTX_L_G | 0x80);
		spi_bb_delay_short();

		for (int i = 0;i < 12;i++) {
			res[i] = spi_bb_exchange_8_mode_3(m_spi_bb, 0);
		}

		spi_bb_end(m_spi_bb);
		chMtxUnlock(&(m_spi_bb->mutex));
		ok = true;
	} else if (m_hwspi_dev) {
		uint8_t txb[1];
		uint8_t rxb[1];

		spiAcquireBus(m_hwspi_dev);
		spiSelect(m_hwspi_dev);
		txb[0] = LSM6DS3_ACC_GYRO_OUTX_L_G | 0x80;
		spiExchange(m_hwspi_dev, 1, txb, rxb);
		spi_bb_delay_short();
		spiExchange(m_hwspi_dev, 12, spi_tx_0_12, res);
		spiUnselect(m_hwspi_dev);
		spiReleaseBus(m_hwspi_dev);
		ok = true;
	}

	return ok;
}

static void terminal_read_reg(int argc, const char **argv) {
	if (argc == 2) {
		int reg = -1;
		sscanf(argv[1], "%d", &reg);

		if (reg >= 0) {
			uint8_t res = 0;
			bool ok = read_reg(reg, &res);

			if (ok) {
				char bl[9];
				utils_byte_to_binary(res & 0xFF, bl);
				commands_printf("Reg: %s", bl);
			} else {
				commands_printf("Read failed\n");
			}
		} else {
			commands_printf("Invalid argument(s)\n");
		}
	} else {
		commands_printf("This command requires one argument\n");
	}
}

static THD_FUNCTION(lsm6ds3_thread, arg) {
	(void)arg;
	chRegSetThreadName("LSM6DS3");

	// One tick less to ensure the frequency is at least the configured one or a bit higher (US2ST rounds up)
	const systime_t interval = US2ST(1000000 / rate_hz) - 1;

	while (!chThdShouldTerminateX()) {
		systime_t start_time = chVTGetSystemTimeX();

		uint8_t rxb[12];
		bool res = read_gyro_accel(rxb);

		if (res) {
			// Parse 6 axis values
			float gx = (float)((int16_t)((uint16_t)rxb[1] << 8) + rxb[0]) * 4.375 * (2000 / 125) / 1000;
			float gy = (float)((int16_t)((uint16_t)rxb[3] << 8) + rxb[2]) * 4.375 * (2000 / 125) / 1000;
			float gz = (float)((int16_t)((uint16_t)rxb[5] << 8) + rxb[4]) * 4.375 * (2000 / 125) / 1000;
			float ax = (float)((int16_t)((uint16_t)rxb[7] << 8) + rxb[6]) * 0.061 * (16 >> 1) / 1000;
			float ay = (float)((int16_t)((uint16_t)rxb[9] << 8) + rxb[8]) * 0.061 * (16 >> 1) / 1000;
			float az = (float)((int16_t)((uint16_t)rxb[11] << 8) + rxb[10]) * 0.061 * (16 >> 1) / 1000;

			if (res && read_callback) {
				float tmp_accel[3] = {ax,ay,az}, tmp_gyro[3] = {gx,gy,gz}, tmp_mag[3] = {1,2,3};
				read_callback(tmp_accel, tmp_gyro, tmp_mag);
			}
		} else {
			chThdSleep(1);
			continue;
		}

		systime_t sleep_ticks = 1;
		systime_t remaining_sleep_time = start_time + interval - chVTGetSystemTimeX();
		if (remaining_sleep_time > 0 && remaining_sleep_time <= interval) {
			sleep_ticks = remaining_sleep_time;
		}

		chThdSleep(sleep_ticks);
	}
}