#include <Wire.h>
#include <SPI.h>
#include <Adafruit_BMP280.h>
#include <Adafruit_Sensor.h>
#include <Adafruit_MPU6050.h>

#define BMP_SCK  (13)
#define BMP_MISO (12)
#define BMP_MOSI (11)
#define BMP_CS   (10)

// Initialize BMP280 sensor (hardware SPI)
Adafruit_BMP280 bmp(BMP_CS);
Adafruit_MPU6050 mpu;

// Kalman filter variables for altitude estimation
double altitude_estimate = 0;
double velocity_estimate = 0;
double P[2][2] = { {1, 0}, {0, 1} };
double Q[2][2] = { {0.01, 0}, {0, 0.02} };
double R = 0.1;
double K[2];

// Rocket and simulation parameters
double rocket_mass = 0.824; // Mass of the rocket in kg
double time_step = 0.01; // Time step for the simulation in seconds
int thrust_index = 0;

// Arrays to store thrust data
const int thrust_data_length = 240;
double thrust_data[thrust_data_length] = {
    1.167,3.5,5.833,33.5,75,101.5,114.5,114.5,106.25,92.75,84.75,82.25,81.833,83.5,85.109,86.141,86.477,86.789,87.084,87.442,87.98,88.786,89.941,91.229,92.341,93.268,94.278,95.379,96.382,97.035,97.249,97.428,97.593,97.753,97.911,98.072,98.247,98.452,98.734,99.156,99.496,99.648,99.776,99.894,100.168,101.188,102.384,103.786,105.716,107.12,99.771,89.663,81.528,74.099,66.836,59.329,49.907,37.665,24.045,14.673,7.246,3.314,0.78

    // Continue for all values in the CSV
};

// Bias for accelerometer Z-axis
double accel_bias = 0;

void setup() {
    Serial.begin(115200);
    Wire.begin();

    if (!bmp.begin(0x76)) {
        Serial.println("Could not find a valid BMP280 sensor, check wiring!");
        while (1);
    }

    if (!mpu.begin()) {
        Serial.println("Failed to find MPU6050 chip");
        while (1);
    }

    Serial.println("Calibrating MPU6050...");
    calibrateMPU6050();
    Serial.print("Accel Bias: ");
    Serial.println(accel_bias);

    // BMP280 settings
    bmp.setSampling(Adafruit_BMP280::MODE_NORMAL,
                    Adafruit_BMP280::SAMPLING_X2,
                    Adafruit_BMP280::SAMPLING_X16,
                    Adafruit_BMP280::FILTER_X16,
                    Adafruit_BMP280::STANDBY_MS_1);
}

void loop() {
    // Read sensors
    sensors_event_t accel, gyro, temp;
    mpu.getEvent(&accel, &gyro, &temp);
    double measured_accelZ = accel.acceleration.z - accel_bias;
    double measured_altitude = bmp.readAltitude(1013.25); // Adjust sea level pressure as needed

    static double simulated_altitude = 0;
    static double simulated_velocity = 0;
    if (thrust_index < thrust_data_length) {
        double thrust = thrust_data[thrust_index];
        double acceleration = thrust / rocket_mass - 9.81; // Subtract gravity
        simulated_velocity += acceleration * time_step;
        simulated_altitude += simulated_velocity * time_step;
        thrust_index++;
    }

    kalmanUpdate(measured_altitude, measured_accelZ);

    Serial.print("Measured Altitude: ");
    Serial.print(measured_altitude);
    Serial.print(", Estimated Altitude: ");
    Serial.println(altitude_estimate);

    delay(100);
}

void calibrateMPU6050() {
    long accelZ_sum = 0;
    int samples = 1000;
    for (int i = 0; i < samples; i++) {
        sensors_event_t accel, gyro, temp;
        mpu.getEvent(&accel, &gyro, &temp);
        accelZ_sum += accel.acceleration.z;
        delay(2);
    }
    accel_bias = accelZ_sum / samples;
}

void kalmanUpdate(double measured_altitude, double measured_accel) {
    // Predict step
    double dt = time_step;
    altitude_estimate += velocity_estimate * dt + 0.5 * measured_accel * dt * dt;
    velocity_estimate += measured_accel * dt;
    
    P[0][0] += P[1][1] * dt * dt + Q[0][0];
    P[0][1] += P[1][1] * dt;
    P[1][0] += P[1][1] * dt;
    P[1][1] += Q[1][1];

    // Update step
    double y = measured_altitude - altitude_estimate;
    double S = P[0][0] + R;
    K[0] = P[0][0] / S;
    K[1] = P[1][0] / S;

    altitude_estimate += K[0] * y;
    velocity_estimate += K[1] * y;

    double P00_temp = P[0][0];
    double P01_temp = P[0][1];

    P[0][0] -= K[0] * P00_temp;
    P[0][1] -= K[0] * P01_temp;
    P[1][0] -= K[1] * P00_temp;
    P[1][1] -= K[1] * P01_temp;
}