#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.01, 0}, {0, 0.02} };
double R = 0.1;
double K[2];

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

// Array to store thrust data from the CSV file
const int thrust_data_length = 63; 
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
};

// Array to store rocket mass data (converted to kg from g)
const int mass_data_length = 63; 
double rocket_mass_data[mass_data_length] = {
  0.824051, 0.824024, 0.823998, 0.823742, 0.823171, 0.822397, 0.821524, 0.820651, 0.819893, 0.819134, 0.818498, 0.817861, 0.817225, 0.816588, 0.815996, 0.815415, 0.81487, 0.814357, 0.81387, 0.813141, 0.812047, 0.810405, 0.808129, 0.806037, 0.80387, 0.801656, 0.799068, 0.79542, 0.791761, 0.790358, 0.789183, 0.78812, 0.787111, 0.786126, 0.785137, 0.784101, 0.782949, 0.781538, 0.779422, 0.776248, 0.775179, 0.774336, 0.773579, 0.77286, 0.772125, 0.77134, 0.770466, 0.769395, 0.767788, 0.765342, 0.764486, 0.763845, 0.763281, 0.762744, 0.762204, 0.761651, 0.761327, 0.76092, 0.760611, 0.760328, 0.760214, 0.760118, 0.760078
};

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

    // Convert mass data from grams to kilograms
    for (int i = 0; i < mass_data_length; i++) {
        rocket_mass_data[i] /= 1000.0;
    }
}

void loop() {
    mpu.update();
    double accelZ = mpu.getAccZ() - accel_bias;

    static double simulated_altitude = 0;
    static double simulated_velocity = 0;
    if (thrust_index < thrust_data_length) {
        double thrust = thrust_data[thrust_index];
        double rocket_mass = rocket_mass_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++;
    }

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