#include <stdio.h>
#include <math.h>
#include "pico/stdlib.h"
#include "hardware/i2c.h"

const float ACCEL_SENS = 4096.0;
const float GYRO_SENS = 32.8; 
const float diff_threshold = 0.4;
const float tau=0.02;

const uint8_t MPU = 0x68; // MPU6050 I2C address
float AccX, AccY, AccZ;
float GyroX, GyroY, GyroZ;
float accAngleX=0, accAngleY=0, gyroAngleX=0, gyroAngleY=0, gyroAngleZ=0;
float roll, pitch, yaw;
float AccErrorX=0, AccErrorY=0, GyroErrorX=0, GyroErrorY=0, GyroErrorZ=0;
float elapsedTime, currentTime, previousTime;
float rpy[3][3]={{0,0,0},{0,0,0},{0,0,0}};
bool firsttime=1;
int c = 0;
int status = 0;
i2c_inst_t* i2c_instance = i2c0;

void i2c_write_reg(uint8_t reg, uint8_t data);
uint16_t i2c_read_reg16(uint8_t reg);
void calculate_IMU_error();
void setup();
void loop();

int main() {
    setup();

    while (1) {
        loop();
    }

    return 0;
}

void i2c_write_reg(uint8_t reg, uint8_t data) {
    uint8_t buf[] = {reg, data};
    i2c_write_blocking(i2c_instance, MPU, buf, 2, false);
}

uint16_t i2c_read_reg16(uint8_t reg) {
    uint8_t buf[2];
    i2c_write_blocking(i2c_instance, MPU, &reg, 1, true);
    i2c_read_blocking(i2c_instance, MPU, buf, 2, false);
    return (buf[0] << 8) | buf[1];
}
void calculate_IMU_error() {

    //todo optimize calibration squence. when a single while is used it drifts a ton. Could use raw data instead of converted version
    int cycle = 1000;
    while (c < cycle) {
        AccX = (int16_t)i2c_read_reg16(0x3B) / ACCEL_SENS;
        AccY = (int16_t)i2c_read_reg16(0x3D) / ACCEL_SENS;
        AccZ = (int16_t)i2c_read_reg16(0x3F) / ACCEL_SENS;
        AccErrorX = AccErrorX + ((atan((AccY) / sqrt(pow((AccX), 2) + pow((AccZ), 2))) * 180 / M_PI));
        AccErrorY = AccErrorY + ((atan(-1 * (AccX) / sqrt(pow((AccY), 2) + pow((AccZ), 2))) * 180 / M_PI));
        c++;
    }

    AccErrorX = AccErrorX / cycle;
    AccErrorY = AccErrorY / cycle;
    c = 0;

    while (c < cycle) {
        GyroX = (int16_t)i2c_read_reg16(0x43) / GYRO_SENS;
        GyroY = (int16_t)i2c_read_reg16(0x45) / GYRO_SENS;
        GyroZ = (int16_t)i2c_read_reg16(0x47) / GYRO_SENS;
        GyroErrorX = GyroErrorX + GyroX;
        GyroErrorY = GyroErrorY + GyroY;
        GyroErrorZ = GyroErrorZ + GyroZ;
        c++;
    }

    GyroErrorX = GyroErrorX / cycle;
    GyroErrorY = GyroErrorY / cycle;
    GyroErrorZ = GyroErrorZ / cycle;

    printf("AccErrorX: %.2f\n", AccErrorX);
    printf("AccErrorY: %.2f\n", AccErrorY);
    printf("GyroErrorX: %.2f\n", GyroErrorX);
    printf("GyroErrorY: %.2f\n", GyroErrorY);
    printf("GyroErrorZ: %.2f\n", GyroErrorZ);
}

int calculate_status(float roll, float pitch, float yaw) {
    

    
    //updates past rpy

    for (int i=0; i<3; i++){
    rpy[2][i] = rpy[1][i];

    }
    
    for (int i=0; i<3; i++){
    rpy[1][i] = rpy[0][i];

    }

    rpy[0][0]= roll;
    rpy[0][1]= pitch;
    rpy[0][2]= yaw;
    

    //compares past and new rpy
    int status = 0;
    for(int i=0; i<3; i++){
        printf("^^%f^^ ",  fabs(rpy[0][i]-rpy[1][i]));
    if (fabs(rpy[0][i]-rpy[1][i])>diff_threshold && fabs(rpy[1][i]-rpy[2][i])>diff_threshold){
        status = 1;
    }
    }
    
return status;
    
}

void setup() {
    stdio_init_all();
    i2c_init(i2c_instance, 400000);
    gpio_set_function(4, GPIO_FUNC_I2C);
    gpio_set_function(5, GPIO_FUNC_I2C);
    gpio_pull_up(4);
    gpio_pull_up(5);
    sleep_ms(2000);

    //values referenced from https://github.com/alex-mous/MPU6050-C-CPP-Library-for-Raspberry-Pi/blob/master/MPU6050.h 
    i2c_write_reg(0x6B, 0x00); //takes mpu6050 out of sleep mode
    i2c_write_reg(0x1a, 0b00000011); //set DLPF to 44hz (digital low pass filter)
    i2c_write_reg(0x19, 0b00000100); //set sample rate divider to 200hz
    i2c_write_reg(0x1c, 0b00010000); //set acc mode to 2 (changes sensitivity to 8G/second)
    i2c_write_reg(0x1b, 0b00010000); //set gyro mode to 2 (changes sensitivity to 1000degrees/second)
    calculate_IMU_error();
    sleep_ms(20);
}



void loop() {

    AccX = (int16_t)i2c_read_reg16(0x3B) / ACCEL_SENS;
    AccY = (int16_t)i2c_read_reg16(0x3D) / ACCEL_SENS;
    AccZ = (int16_t)i2c_read_reg16(0x3F) / ACCEL_SENS;
    accAngleX = (atan(AccY / sqrt(pow(AccX, 2) + pow(AccZ, 2))) * 180 / M_PI) - AccErrorX;
    //accAngleY = (atan(-1 * AccX / sqrt(pow(AccY, 2) + pow(AccZ, 2))) * 180 / M_PI) + AccErrorY;
    accAngleY = (atan(-1 * AccX / sqrt(pow(AccY, 2) + pow(AccZ, 2))) * 180 / M_PI) - AccErrorY; //this probably the correct way

    previousTime = currentTime;
    currentTime = to_ms_since_boot(get_absolute_time());
    elapsedTime = (currentTime - previousTime) / 1000;

    GyroX = (int16_t)i2c_read_reg16(0x43) / GYRO_SENS;
    GyroY = (int16_t)i2c_read_reg16(0x45) / GYRO_SENS;
    GyroZ = (int16_t)i2c_read_reg16(0x47) / GYRO_SENS;

    GyroX = GyroX - GyroErrorX;
    GyroY = GyroY - GyroErrorY;
    GyroZ = GyroZ - GyroErrorZ;

    gyroAngleX = gyroAngleX + GyroX * elapsedTime;
    gyroAngleY = gyroAngleY + GyroY * elapsedTime;
    yaw = yaw + GyroZ * elapsedTime;


    roll = (1-tau) * gyroAngleX + (tau) * accAngleX;
    pitch = (1-tau) * gyroAngleY + (tau) * accAngleY;
    
    if (firsttime){
        calculate_status (roll, pitch, yaw);
        calculate_status (roll, pitch, yaw); //init both past rpy values to current
        firsttime = 0;
    }
    else if (status == 0){
        
        status = calculate_status (roll, pitch, yaw);
    } 



    printf("Roll: %.2f / Pitch: %.2f / Yaw: %.2f / Status: %d\n", roll, pitch, yaw, status);

    sleep_ms(10);
}