#include "main.h"
#include <stdio.h>
#include <string.h>
#include "mpu6050.h"
#include <math.h>
#include "madgwick.h"
// FreeRTOS Header
#include "FreeRTOS.h"
#include "task.h"
#include "semphr.h"


UART_HandleTypeDef *huart_mpu;
MPU6050_HandleTypeDef *hmpu;
uint8_t sensorData[14];
TaskHandle_t xSensorReadingTask = NULL;

// using volatile variables so that the calculation of the DMA callback can be handled in the main loop
volatile int16_t raw_data[6];

void sendQuat (UART_HandleTypeDef *uart, MPU6050_HandleTypeDef *mpu, MadgwickFilter *quat) {
	QuaternionData FrameQuat;

	FrameQuat.StartByte = PACKET_START_BYTE;
	FrameQuat.SensorAddress = mpu->address;
	FrameQuat.qw = quat->q[0];
	FrameQuat.qx = quat->q[1];
	FrameQuat.qy = quat->q[2];
	FrameQuat.qz = quat->q[3];
	FrameQuat.EndByte = PACKET_END_BYTE;

	HAL_UART_Transmit(uart, (uint8_t*)&FrameQuat, sizeof(FrameQuat), 50);
}

void read_dma (MPU6050_HandleTypeDef *mpu, uint8_t binaryData[]) {
	  if(HAL_I2C_Mem_Read_DMA(mpu->hi2c, mpu->address, MPU6050_ACC_MEAS_REG, I2C_MEMADD_SIZE_8BIT, binaryData, 14) == HAL_OK) {
		  ulTaskNotifyTake(pdTRUE, pdMS_TO_TICKS(10));
	  }
}

float* process_binaryData (MPU6050_HandleTypeDef *mpu, uint8_t binarySensorData[], float processedData[]) {
	int16_t rawData[6]; //[Ax, Ay, Az, Gx, Gy, Gz]

	// accelerometer
    rawData[0] = (int16_t)(binarySensorData[0] << 8 | binarySensorData[1]);
    rawData[1] = (int16_t)(binarySensorData[2] << 8 | binarySensorData[3]);
    rawData[2] = (int16_t)(binarySensorData[4] << 8 | binarySensorData[5]);
	// gyro
    rawData[3] = (int16_t)(binarySensorData[8] << 8 | binarySensorData[9]);
    rawData[4] = (int16_t)(binarySensorData[10] << 8 | binarySensorData[11]);
    rawData[5] = (int16_t)(binarySensorData[12] << 8 | binarySensorData[13]);

    // processedData [ax, ay, az, gx, gy, gz]
	processedData[0] = (float)rawData[0]/mpu->accel_factor;
	processedData[1] = (float)rawData[1]/mpu->accel_factor;
	processedData[2] = (float)rawData[2]/mpu->accel_factor;
	processedData[3] = (float)rawData[3]/mpu->gyro_factor;
	processedData[4] = (float)rawData[4]/mpu->gyro_factor;
	processedData[5] = (float)rawData[5]/mpu->gyro_factor;

	if (mpu->calibrate) {
		processedData[0] -=  mpu->accel_bias[0];
		processedData[1] -=  mpu->accel_bias[1];
		processedData[2] -=  mpu->accel_bias[2];
		processedData[3] -=  mpu->gyro_bias[0];
		processedData[4] -=  mpu->gyro_bias[1];
		processedData[5] -=  mpu->gyro_bias[2];
	}
	return processedData;
}

void HAL_I2C_MemRxCpltCallback(I2C_HandleTypeDef *I2C_address) {
	// all this function does now is alarm the sensor reading task that the
	// dma read is finished
	BaseType_t xHigherPriorityTaskWoken = pdFALSE;
	if (I2C_address -> Instance == I2C1) {
		vTaskNotifyGiveFromISR(xSensorReadingTask, &xHigherPriorityTaskWoken);
	}
	portYIELD_FROM_ISR(xHigherPriorityTaskWoken);

}


void Init_MPU6050 (MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart) {
    uint8_t data;
    uint8_t readback;
    HAL_StatusTypeDef status;
    char buffer[128];
    huart_mpu = uart;
    hmpu = mpu;

    sprintf(buffer, "Initializing MPU6050 at address: 0x%02X\n", mpu->address >> 1);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 25);

    status = HAL_I2C_IsDeviceReady(mpu->hi2c, mpu->address, 15, 25);

    if (status != HAL_OK) {
        sprintf(buffer, "Error: Device not ready! Status: %d\n", status);
        HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 25);
    }

    // Check WHO_AM_I register
    status = HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, 0x75, I2C_MEMADD_SIZE_8BIT, &readback, 1, 25);
    sprintf(buffer, "WHO_AM_I register: 0x%02X (should be 0x68)\n", readback);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 25);

    if (readback != 0x68) {
        sprintf(buffer, "Warning: Unexpected device ID!\n");
        HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 25);
    }

    // Reset the device
    data = 0x80;  // Reset bit
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_PWR_MNG_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);
    HAL_Delay(100);  // Wait for reset to complete

    // Power Management - set clock source
    data = 0x01;  // PLL with X-axis gyro reference
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_PWR_MNG_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);

    // Read back to verify
    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_PWR_MNG_REG, I2C_MEMADD_SIZE_8BIT, &readback, 1, 250);
    sprintf(buffer, "PWR_MGMT_1: Wrote 0x01, Read 0x%02X\n", readback);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    // Configure DLPF (Digital Low Pass Filter)
    data = 0x03;  // ~44Hz bandwidth
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_CONF_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);

    // Sample Rate Divider
    data = 0x07;  // 1kHz / (1 + 7) = 125Hz
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_SMP_RATE_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);

    // Set conversion factors based on range settings
    setRange_MPU(mpu, uart);

    // Gyro Config
    data = (mpu->gyro_range << 3);  // Shift left 3 bits
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_GYR_CONF_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);
    HAL_Delay(10);

    // Read back to verify
    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_GYR_CONF_REG, I2C_MEMADD_SIZE_8BIT, &readback, 1, 250);
    sprintf(buffer, "GYRO_CONFIG: Wrote 0x%02X, Read 0x%02X\n", (mpu->gyro_range << 3), readback);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    // Accel Config
    data = (mpu->accel_range << 3);  // Shift left 3 bits
    HAL_I2C_Mem_Write(mpu->hi2c, mpu->address, MPU6050_ACC_CONF_REG, I2C_MEMADD_SIZE_8BIT, &data, 1, 250);
    HAL_Delay(10);

    // Read back to verify
    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_ACC_CONF_REG, I2C_MEMADD_SIZE_8BIT, &readback, 1, 250);
    sprintf(buffer, "ACCEL_CONFIG: Wrote 0x%02X, Read 0x%02X\n", (mpu->accel_range << 3), readback);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    // Print current configuration
    sprintf(buffer, "Configuration complete. Accel factor: %.2f, Gyro factor: %.2f\n",
            mpu->accel_factor, mpu->gyro_factor);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);
    HAL_Delay(100);

    // Start first measurement using DMA
    HAL_I2C_Mem_Read_DMA(mpu->hi2c, mpu->address, MPU6050_ACC_MEAS_REG, I2C_MEMADD_SIZE_8BIT, sensorData, 14);
	}

void setRange_MPU (MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart) {
	char buffer[128];

	switch (mpu->accel_range) {
	case 0x00:
		mpu->accel_factor = 16384.0;
		break;
	case 0x01:
		mpu->accel_factor = 8192.0;
		break;
	case 0x02:
		mpu->accel_factor = 4096.0;
		break;
	case 0x03:
		mpu->accel_factor = 2048.0;
		break;
	default:
		mpu->accel_factor = 16384.0;
		sprintf(buffer, "Invalid range provided for Accelerometer: %d. Setting factor to default %f\n",mpu->accel_range, mpu->accel_factor);
		HAL_UART_Transmit(uart, (uint8_t *)buffer, strlen(buffer), 100);
	}

	switch (mpu->gyro_range) {
	case 0x00:
		mpu->gyro_factor = 131.0;
		break;
	case 0x01:
		mpu->gyro_factor = 65.5;
		break;
	case 0x02:
		mpu->gyro_factor = 32.8;
		break;
	case 0x03:
		mpu->gyro_factor = 16.4;
		break;
	default:
		mpu->accel_factor = 131.0;
		sprintf(buffer, "Invalid range provided for Gyroscope: %d. Setting factor to default %f\n", mpu->gyro_range, mpu->gyro_factor);
		HAL_UART_Transmit(uart, (uint8_t *)buffer, strlen(buffer), 100);
	}
}

void calibrate_Sensor(MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart) {
	const uint16_t num_samples = 1500;
    uint16_t samples_taken = 0;
    uint8_t binarySensorData[14];
    float processedData[6];
    float accel_sum[3] = {0.0, 0.0, 0.0};
    float gyro_sum[3] = {0.0, 0.0, 0.0};
    char buffer[256];

    sprintf(buffer, "Starting calibration at address: 0x%02X\nKeep the sensor still during calibration...\n", mpu->address >> 1);
    //sprintf(buffer + strlen(buffer), "Keep the sensor still during calibration...\n");
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    for (int i = 0; i < num_samples; i++) {
		read_dma(mpu, binarySensorData);
		process_binaryData(mpu, binarySensorData, processedData);
        accel_sum[0] += processedData[0];
        accel_sum[1] += processedData[1];
        accel_sum[2] += processedData[2];
        gyro_sum[0] += processedData[3];
        gyro_sum[1] += processedData[4];
        gyro_sum[2] += processedData[5];

        // Give progress to user every 100 samples taken
        if (i % 100 == 0) {
            sprintf(buffer, "Calibration progress: %d%%\n", (i*100)/num_samples);
            HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);
        }

        samples_taken++;

    }

    if (samples_taken > 0) {
        // Calculate bias values
        mpu->accel_bias[0] = accel_sum[0] / samples_taken;
        mpu->accel_bias[1] = accel_sum[1] / samples_taken;
        mpu->accel_bias[2] = accel_sum[2] / samples_taken;

        // Subtract 1g from z-axis to account for gravity
        mpu->accel_bias[2] -= 1.0f;

        mpu->gyro_bias[0] = gyro_sum[0] / samples_taken;
        mpu->gyro_bias[1] = gyro_sum[1] / samples_taken;
        mpu->gyro_bias[2] = gyro_sum[2] / samples_taken;

        // Set the calibration flag
        mpu->calibrate = 1;

        // Print results
        sprintf(buffer, "Calibration successful!\n");
        sprintf(buffer + strlen(buffer), "Accelerometer bias (g): ");
        sprintf(buffer + strlen(buffer), "X: %.4f, Y: %.4f, Z: %.4f\n",
                mpu->accel_bias[0], mpu->accel_bias[1], mpu->accel_bias[2]);
        sprintf(buffer + strlen(buffer), "Gyroscope bias (°/s): ");
        sprintf(buffer + strlen(buffer), "X: %.4f, Y: %.4f, Z: %.4f\n",
                mpu->gyro_bias[0], mpu->gyro_bias[1], mpu->gyro_bias[2]);
        HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);
    } else {
        sprintf(buffer, "Calibration failed: No samples collected\n");
        HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);
    }
}

// custom function to read device status, because HAL_I2C_IsDeviceReady is apparently
// often buggy in its implementation - not using it currently
HAL_StatusTypeDef check_mpu6050_ready(MPU6050_HandleTypeDef *mpu) {
    uint8_t dummy;
    return HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, 0x75, I2C_MEMADD_SIZE_8BIT, &dummy, 1, 100);
}



// Debug functions
void MPU_Diagnostics(MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart) {
    char buffer[256];
    uint8_t reg_value;

    // Print MPU settings and status
    sprintf(buffer, "MPU6050 Diagnostics:\n----------------------\n");
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Address: 0x%02X\n", mpu->address >> 1);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Calibration status: %s\n", mpu->calibrate ? "Calibrated" : "Not calibrated");
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Accel range setting: 0x%02X\n", mpu->accel_range);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Gyro range setting: 0x%02X\n", mpu->gyro_range);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Accel factor: %.2f\n", mpu->accel_factor);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Gyro factor: %.2f\n", mpu->gyro_factor);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Accel bias: X=%.4f Y=%.4f Z=%.4f\n",
            mpu->accel_bias[0], mpu->accel_bias[1], mpu->accel_bias[2]);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "Gyro bias: X=%.4f Y=%.4f Z=%.4f\n",
            mpu->gyro_bias[0], mpu->gyro_bias[1], mpu->gyro_bias[2]);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    // Read WHO_AM_I register
    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, 0x75, I2C_MEMADD_SIZE_8BIT, &reg_value, 1, 100);
    sprintf(buffer, "WHO_AM_I (0x75): 0x%02X (should be 0x68)\n", reg_value);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    // Read register values to verify configuration
    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_PWR_MNG_REG, I2C_MEMADD_SIZE_8BIT, &reg_value, 1, 100);
    sprintf(buffer, "PWR_MGMT_1 (0x6B): 0x%02X\n", reg_value);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_GYR_CONF_REG, I2C_MEMADD_SIZE_8BIT, &reg_value, 1, 100);
    sprintf(buffer, "GYRO_CONFIG (0x1B): 0x%02X\n", reg_value);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_ACC_CONF_REG, I2C_MEMADD_SIZE_8BIT, &reg_value, 1, 100);
    sprintf(buffer, "ACCEL_CONFIG (0x1C): 0x%02X\n", reg_value);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);

    sprintf(buffer, "----------------------\n");
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 100);
}

void checkRegisters(MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart){
	char buffer[128];
	uint8_t answer;
    huart_mpu = uart;
    hmpu = mpu;

	sprintf(buffer, "-------------------------------\n");
	HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_GYR_CONF_REG, I2C_MEMADD_SIZE_8BIT, &answer, 1, 250);
    sprintf(buffer, "Gyro Config: 0x%02X\n", answer);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_ACC_CONF_REG, I2C_MEMADD_SIZE_8BIT, &answer, 1, 250);
    sprintf(buffer, "ACC Config: 0x%02X\n", answer);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

    HAL_I2C_Mem_Read(mpu->hi2c, mpu->address, MPU6050_PWR_MNG_REG, I2C_MEMADD_SIZE_8BIT, &answer, 1, 250);
    sprintf(buffer, "PWR MNG Config: 0x%02X\n", answer);
    HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);

	sprintf(buffer, "-------------------------------\n\n\n\n");
	HAL_UART_Transmit(uart, (uint8_t*)buffer, strlen(buffer), 250);
}

void I2C_Scanner(I2C_HandleTypeDef *hi2c, UART_HandleTypeDef *huart) {
    char buffer[64];
    uint8_t found = 0;

    for (uint8_t addr = 0x08; addr < 0x78; addr++) {
        if (HAL_I2C_IsDeviceReady(hi2c, addr << 1, 1, 10) == HAL_OK) {
            sprintf(buffer, "Device found: 0x%02X\n", addr);
            HAL_UART_Transmit(huart, (uint8_t*)buffer, strlen(buffer), 100);
            found = 1;
        }
    }

    if (!found) {
        sprintf(buffer, "No devices found\n");
        HAL_UART_Transmit(huart, (uint8_t*)buffer, strlen(buffer), 100);
    }
}
