#include "main.h"
#include <stdio.h>
#include <string.h>
#include "mpu6050.h"
#include <math.h>


UART_HandleTypeDef *huart_mpu;
MPU6050_HandleTypeDef *hmpu;
uint8_t sensorData[14];

// using volatile variables so that the calculation of the DMA callback can be handled in the main loop
volatile int16_t raw_data[6];
volatile uint8_t data_ready = 0;
volatile uint8_t dma_ready = 1;


// new function to initiate dma transter
HAL_StatusTypeDef start_dma_transfer(MPU6050_HandleTypeDef *mpu) {
	// check if previous transfer is complete before initiating new one, else return HAL_BUSY
	if (dma_ready) {
		dma_ready = 0;
		return HAL_I2C_Mem_Read_DMA(mpu->hi2c, mpu->address, MPU6050_ACC_MEAS_REG, I2C_MEMADD_SIZE_8BIT, sensorData, 14);
	}
	return HAL_BUSY;
}

// 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);
}

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_MPU (MPU6050_HandleTypeDef *mpu, UART_HandleTypeDef *uart){
    const uint16_t num_samples = 1000;
    int32_t accel_sum[3] = {0, 0, 0};
    int32_t gyro_sum[3] = {0, 0, 0};
    int16_t local_data[6];
    char buffer[256];
    uint16_t i = 0, samples_taken = 0;
    hmpu=mpu;

    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);

    // check if dma ready and i2c
    data_ready = 0;
    dma_ready = 1;

    if (HAL_I2C_GetState(mpu->hi2c) == HAL_I2C_STATE_READY) {
    	start_dma_transfer(hmpu);
    }

    // Check if data is already being collected, start otherwise
//    if (!data_ready) {
//        HAL_I2C_Mem_Read_DMA(mpu->hi2c, mpu->address, MPU6050_ACC_MEAS_REG, I2C_MEMADD_SIZE_8BIT, sensorData, 14);
//        HAL_Delay(10);
//    }

    // Discard first 50 readings to let the sensor stabilize
    for (i = 0; i < 50;) {
        if (data_ready) {
            data_ready = 0;
            i++;

            start_dma_transfer(hmpu);
        }
        HAL_Delay(1);
    }

    // Reset counters for actual calibration
    i = 0;
    samples_taken = 0;

    // Collect and sum the samples
    while (i < num_samples) {
        if (data_ready) {
            // Copy sensor data to local array to avoid data loss
            local_data[0] = raw_data[0];
            local_data[1] = raw_data[1];
            local_data[2] = raw_data[2];
            local_data[3] = raw_data[3];
            local_data[4] = raw_data[4];
            local_data[5] = raw_data[5];

            // Reset data gathering flag
            data_ready = 0;

            // Add to the sum arrays
            accel_sum[0] += local_data[0];
            accel_sum[1] += local_data[1];
            accel_sum[2] += local_data[2];
            gyro_sum[0] += local_data[3];
            gyro_sum[1] += local_data[4];
            gyro_sum[2] += local_data[5];

            samples_taken++;
            i++;

            // 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);
            }

            start_dma_transfer(hmpu);
        } else {
            // Small delay to prevent CPU maxing out
            HAL_Delay(1);
        }
    }

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

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

        mpu->gyro_bias[0] = (float)gyro_sum[0] / samples_taken / mpu->gyro_factor;
        mpu->gyro_bias[1] = (float)gyro_sum[1] / samples_taken / mpu->gyro_factor;
        mpu->gyro_bias[2] = (float)gyro_sum[2] / samples_taken / mpu->gyro_factor;

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

        // Print results
        sprintf(buffer, "Calibration successful!\n");
        sprintf(buffer + strlen(buffer), "Accelerometer bias (g):\n");
        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):\n");
        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);
    }
}





void HAL_I2C_MemRxCpltCallback(I2C_HandleTypeDef *I2C_address) {
	if (I2C_address -> Instance == I2C1) {
		// copy data from sensor to array
		// accelerometer
		raw_data[0] = (int16_t)(sensorData[0] << 8 | sensorData[1]);
		raw_data[1] = (int16_t)(sensorData[2] << 8 | sensorData[3]);
		raw_data[2] = (int16_t)(sensorData[4] << 8 | sensorData[5]);
		// gyro
		raw_data[3] = (int16_t)(sensorData[8] << 8 | sensorData[9]);
		raw_data[4] = (int16_t)(sensorData[10] << 8 | sensorData[11]);
		raw_data[5] = (int16_t)(sensorData[12] << 8 | sensorData[13]);

		// set the flag for the main loop
		data_ready = 1;
		dma_ready = 1;

		// Start next measurement
		// moved to main for the time being
		//HAL_I2C_Mem_Read_DMA(I2C_address, hmpu->address, MPU6050_ACC_MEAS_REG, I2C_MEMADD_SIZE_8BIT, sensorData, 14);
	}
}


// 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);
}
