// Copyright (c) 2023 Oleg Kalachev <okalachev@gmail.com>
// Repository: https://github.com/okalachev/flix

// Work with the IMU sensor

#include <SPI.h>
#include <Wire.h>
#include <FlixPeriph.h>
#include "vector.h"
#include "filter.h"
#include "util.h"

IMU *imu;
int imuModel = -1; // 1 - MPU9250, 2 - ICM20948, 3 - MPU6050, 4 - ICM40609D
int imuBus = 0; // 0 - SPI, 1 - I2C
int imuSckPin = SCK, imuMisoPin = MISO, imuMosiPin = MOSI, imuCsPin = SS, imuIntPin = -1;
int imuSdaPin = SDA, imuSclPin = SCL;
Vector imuRotation(0, 0, PI / 2); // imu orientation as Euler angles

Vector gyro; // gyroscope output, rad/s
Vector gyroBias;

Vector acc; // accelerometer output, m/s/s
Vector accBias;
Vector accScale(1, 1, 1);

LowPassFilter<Vector> gyroBiasFilter(0.001);

void setupIMU() {
	print("Setup IMU\n");
	free(imu);
	if (imuModel == 3) imuBus = 1; // MPU6050 is I2C only

	if (imuBus == 0) {
		// SPI connection
#if defined(ESP32)
		SPI.begin(imuSckPin, imuMisoPin, imuMosiPin);
#elif defined(ARDUINO_ARCH_STM32)
		SPI.setSCLK(imuSckPin);
		SPI.setMOSI(imuMosiPin);
		SPI.setMISO(imuMisoPin);
		SPI.begin();
#endif
		imu = IMU::create(imuModel, SPI, imuCsPin, imuIntPin);
	} else {
		// I2C connection
#if defined(ESP32)
		Wire.setPins(imuSdaPin, imuSclPin);
#elif defined(ARDUINO_ARCH_STM32)
		Wire.setSDA(imuSdaPin);
		Wire.setSCL(imuSclPin);
#endif
		imu = IMU::create(imuModel, Wire, imuIntPin);
	}

	imu->begin();
	configureIMU();
}

void configureIMU() {
	imu->setAccelRange(IMU::ACCEL_RANGE_4G);
	imu->setGyroRange(IMU::GYRO_RANGE_2000DPS);
	imu->setDLPF(IMU::DLPF_MAX);
	imu->setRate(IMU::RATE_1KHZ_APPROX);
	imu->setupInterrupt();
}

void readIMU() {
	imu->waitForData();
	imu->getGyro(gyro.x, gyro.y, gyro.z);
	imu->getAccel(acc.x, acc.y, acc.z);
	calibrateGyroOnce();

	// Apply scale and bias
	acc = (acc - accBias) / accScale;
	gyro = gyro - gyroBias;

	// Rotate to body frame
	Quaternion rotation = Quaternion::fromEuler(imuRotation);
	acc = Quaternion::rotateVector(acc, rotation.inversed());
	gyro = Quaternion::rotateVector(gyro, rotation.inversed());
}

void calibrateGyroOnce() {
	static Delay landedDelay(2);
	if (!landedDelay.update(landed)) return; // calibrate only if definitely stationary

	gyroBias = gyroBiasFilter.update(gyro);
}

void calibrateAccel() {
	print("Calibrating accelerometer\n");
	imu->setAccelRange(IMU::ACCEL_RANGE_2G); // the most sensitive mode

	print("1/6 Place level [8 sec]\n");
	pause(8);
	calibrateAccelOnce();
	print("2/6 Place nose up [8 sec]\n");
	pause(8);
	calibrateAccelOnce();
	print("3/6 Place nose down [8 sec]\n");
	pause(8);
	calibrateAccelOnce();
	print("4/6 Place on right side [8 sec]\n");
	pause(8);
	calibrateAccelOnce();
	print("5/6 Place on left side [8 sec]\n");
	pause(8);
	calibrateAccelOnce();
	print("6/6 Place upside down [8 sec]\n");
	pause(8);
	calibrateAccelOnce();

	printIMUCalibration();
	print("✓ Calibration done!\n");
	configureIMU();
}

void calibrateAccelOnce() {
	const int samples = 1000;
	static Vector accMax(-INFINITY, -INFINITY, -INFINITY);
	static Vector accMin(INFINITY, INFINITY, INFINITY);

	// Compute the average of the accelerometer readings
	acc = Vector(0, 0, 0);
	for (int i = 0; i < samples; i++) {
		imu->waitForData();
		Vector sample;
		imu->getAccel(sample.x, sample.y, sample.z);
		acc = acc + sample;
	}
	acc = acc / samples;

	// Update the maximum and minimum values
	if (acc.x > accMax.x) accMax.x = acc.x;
	if (acc.y > accMax.y) accMax.y = acc.y;
	if (acc.z > accMax.z) accMax.z = acc.z;
	if (acc.x < accMin.x) accMin.x = acc.x;
	if (acc.y < accMin.y) accMin.y = acc.y;
	if (acc.z < accMin.z) accMin.z = acc.z;

	// Compute scale and bias
	accScale = (accMax - accMin) / 2 / ONE_G;
	accBias = (accMax + accMin) / 2;
}

void printIMUCalibration() {
	print("gyro bias: %f %f %f\n", gyroBias.x, gyroBias.y, gyroBias.z);
	print("accel bias: %f %f %f\n", accBias.x, accBias.y, accBias.z);
	print("accel scale: %f %f %f\n", accScale.x, accScale.y, accScale.z);
}

void printIMUInfo() {
	imu->status() ? print("status: ERROR %d\n", imu->status()) : print("status: OK\n");
	print("model: %s\n", imu->getModel());
	print("who am I: 0x%02X\n", imu->whoAmI());
	print("rate: %.0f\n", loopRate);
	print("interrupt mode: %s\n", imuIntPin != -1 ? "pin" : "timer");
	print("temperature: %.1f °C\n", imu->getTemp());
	print("gyro: %f %f %f\n", gyro.x, gyro.y, gyro.z);
	print("acc: %f %f %f\n", acc.x, acc.y, acc.z);
	imu->waitForData();
	Vector rawGyro, rawAcc;
	imu->getGyro(rawGyro.x, rawGyro.y, rawGyro.z);
	imu->getAccel(rawAcc.x, rawAcc.y, rawAcc.z);
	print("raw gyro: %f %f %f\n", rawGyro.x, rawGyro.y, rawGyro.z);
	print("raw acc: %f %f %f\n", rawAcc.x, rawAcc.y, rawAcc.z);
}