#include <Wire.h>
#include <SPI.h>
#include <Adafruit_BMP280.h>
#include <Adafruit_Sensor.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);

// Mock MPU6050 class
class MPU6050 {
public:
    void initialize() {
        //Serial.println("MPU6050 initialized.");
    }
    bool testConnection() {
        return true;
    }
    void update() {}
    double getAccZ() {
        static double accel = 9.81; // Simulate constant gravity
        return accel;
    }
};

MPU6050 mpu;
double accel_bias = 0;

// 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.05, 0}, {0, 0.02} };
double R = 0.05;
double K[2];

// Rocket and simulation parameters
double time_step = 0.01; // Time step for the simulation in seconds
int data_index = 0;

// Array to store acceleration data from the image
const int data_length = 63;
double acceleration_data[data_length] = {
    -8.376, -5.545, -2.713, 30.876, 81.318, 113.62, 129.564, 129.698, 119.749, 103.373, 
    93.673, 90.682, 90.235, 92.338, 94.367, 95.689, 96.155, 96.588, 96.997, 97.506, 
    98.269, 99.409, 101.025, 102.776, 104.3, 105.569, 106.903, 108.264, 109.393, 
    110.154, 110.395, 110.608, 110.817, 111.032, 111.261, 111.51, 111.793, 112.138, 
    112.609, 113.212, 113.646, 113.819, 113.953, 114.066, 114.374, 115.636, 117.104, 
    118.799, 121.074, 122.51, 112.724, 99.43, 88.737, 78.987, 69.47, 59.644, 47.303, 
    31.296, 13.511, 1.308, -8.398, -13.528, -16.87
};

// Array to store set altitude data
double altitude_data[data_length] = {
  0, 0, 0, 0.002, 0.009, 0.026, 0.055, 0.097, 0.151, 0.217, 0.292, 0.377, 0.471, 
  0.573, 0.678, 0.781, 0.884, 0.987, 1.089, 1.253, 1.518, 1.965, 2.682, 3.443, 
  4.294, 5.259, 6.516, 8.48, 10.717, 11.63, 12.419, 13.156, 13.875, 14.598, 15.342, 
  16.142, 17.058, 18.214, 20.023, 22.904, 23.92, 24.738, 25.484, 26.203, 26.926, 
  27.693, 28.559, 29.64, 31.303, 33.891, 35.108, 36.033, 36.855, 37.644, 38.442, 
  39.295, 40.268, 41.493, 43.336, 46.106, 47.223, 48.162, 48.558
};

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

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

    mpu.initialize();
    if (!mpu.testConnection()) {
        Serial.println("MPU6050 connection failed!");
        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() {
    mpu.update();
    double accelZ = mpu.getAccZ() - accel_bias;

    double simulated_altitude = 0;
    double simulated_velocity = 0;
    if (data_index < data_length) {
        double acceleration = acceleration_data[data_index];
        simulated_velocity += acceleration * time_step;
        simulated_altitude = altitude_data[data_index]; // Use set altitude data
        data_index++;
    }

    double measured_altitude = simulated_altitude;
    double measured_accel = accelZ;

    kalmanUpdate(measured_altitude, measured_accel);

   // Serial.print("Simulated Altitude: ");
   // Serial.print(simulated_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++) {
        mpu.update();
        accelZ_sum += mpu.getAccZ();
        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;
}