#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 accelZ = 9.81; // Simulate constant gravity
        return accelZ;
    }
    double getAccX() {
        static double accelX = 0.0; // Simulate lateral acceleration
        return accelX;
    }
};

MPU6050 mpu;
double accel_bias_z = 0;
double accel_bias_x = 0;

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

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

// Array to store acceleration data from the CSV file
const int data_length = 269; // Adjust based on the CSV data length
double acceleration_data[data_length] = {
    -4.978,37.158,79.941,90.829,91.48,81.268,77.38,80.431,82.58,83.824,85.069,85.845,86.15,86.455,86.751,87.028,87.284,87.525,87.824,88.246,88.001,87.175,86.039,84.737,83.341,81.853,80.292,78.688,77.474,76.519,75.346,74.198,73.008,71.663,69.815,67.361,65.391,64.155,63.047,61.95,60.763,59.172,56.787,55.011,54.406,53.884,53.345,52.7,51.766,50.777,50.055,49.45,48.755,47.753,46.399,45.426,44.722,43.89,42.696,41.785,41.223,40.965,41.014,41.118,39.896,34.912,29.119,21.027,13.395,9.165,5.718,1.299,-4.712,-9.2,-11.862,-14.677,-18.309,-20.549,-20.536,-20.503,-20.457,-20.391,-20.307,-20.241,-20.181,-20.106,-20.014,-19.94,-19.834,-19.672,-19.439,-19.333,-19.27,-19.218,-19.178,-19.127,-19.072,-19.005,-18.913,-18.779,-18.578,-18.407,-18.339,-18.287,-18.244,-18.191,-18.11,-18.042,-17.948,-17.803,-17.601,-17.48,-17.418,-17.368,-17.313,-17.195,-17.098,-16.932,-16.749,-16.616,-16.544,-16.468,-16.382,-16.312,-16.219,-16.064,-15.891,-15.804,-15.74,-15.666,-15.595,-15.504,-15.424,-15.287,-15.14,-15.004,-14.94,-14.886,-14.818,-14.753,-14.672,-14.567,-14.423,-14.299,-14.196,-14.128,-14.068,-14.002,-13.913,-13.813,-13.697,-13.598,-13.518,-13.458,-13.386,-13.325,-13.245,-13.147,-13.066,-12.996,-12.932,-12.863,-12.804,-12.732,-12.652,-12.578,-12.508,-12.445,-12.379,-12.331,-12.287,-12.237,-12.189,-12.126,-12.054,-11.978,-11.915,-11.861,-11.818,-11.781,-11.742,-11.689,-11.645,-11.606,-11.546,-11.491,-11.44,-11.391,-11.339,-11.294,-11.264,-11.244,-11.21,-11.167,-11.125,-11.082,-11.04,-10.994,-10.96,-10.916,-10.878,-10.841,-10.817,-10.778,-10.754,-10.728,-10.697,-10.665,-10.637,-10.614,-10.588,-10.559,-10.529,-10.505,-10.464,-10.433,-10.415,-10.384,-10.357,-10.331,-10.292,-10.266,-10.269,-10.257,-10.249,-10.236,-10.21,-10.187,-10.173,-10.156,-10.132,-10.096,-10.063,-10.048,-10.025,-10.004,-9.978,-9.961,-9.952,-9.937,-9.928,-9.923,-9.914,-9.916,-9.908,-9.901,-9.886,-9.874,-9.867,-9.864,-9.861,-9.846,-9.83,-9.818,-9.8,-9.785,-9.771,-9.754,-9.728,-9.714,-9.684,-9.649,-9.623,-9.604,-9.571,-9.545,-9.521,-9.485,-9.43    

};

// Array to store set altitude data from the CSV file
double altitude_data[data_length] = {
    0,0.002,0.01,0.026,0.051,0.085,0.127,0.177,0.235,0.301,0.376,0.459,0.551,0.652,0.755,0.859,0.962,1.065,1.229,1.497,1.951,2.747,3.884,5.236,6.842,8.654,10.669,12.883,14.234,15.382,16.5,17.669,18.972,20.584,23.111,26.212,27.831,29.195,30.522,31.877,33.421,35.717,39.272,41.009,42.463,43.932,45.59,47.796,51.18,53.864,55.62,57.38,59.499,62.731,67.277,69.353,71.859,75.053,79.896,81.874,83.666,85.673,88.713,93.344,95.406,97.45,100.078,104.048,107.034,108.994,111.065,114.156,118.789,120.894,122.892,125.29,128.87,129.134,129.529,130.122,131.01,132.339,134.306,136.228,138.621,142.192,145.634,147.5,149.442,152.057,155.955,158.068,159.771,161.566,163.912,167.412,169.814,171.55,173.312,175.547,178.881,182.418,184.199,186.021,188.338,191.793,195.095,197.132,199.141,201.684,205.473,208.261,210.273,212.629,216.14,220.327,222.797,226.35,230.659,234.198,236.788,240.444,244.615,247.503,249.67,252.807,256.857,259.455,261.96,265.417,269.343,273.229,275.545,278.464,282.26,286.018,288.429,291.333,295,298.63,301.303,304.439,307.973,311.471,314.933,318.039,321.434,324.793,328.117,331.407,334.662,337.883,341.069,344.222,347.342,350.428,353.48,356.5,359.486,362.44,365.361,368.251,371.108,373.932,376.726,379.487,382.217,384.916,387.583,390.22,392.826,395.401,397.946,400.46,402.944,405.397,407.821,410.215,412.579,414.914,417.219,419.495,421.741,423.959,426.148,428.307,430.438,432.541,434.614,436.66,438.677,440.666,442.627,444.56,446.465,448.343,450.192,452.014,453.809,455.576,457.316,459.028,460.714,462.372,464.004,465.608,467.186,468.737,470.261,471.759,473.231,474.675,476.094,477.486,478.852,480.191,481.505,482.792,484.054,485.29,486.499,487.684,488.842,489.975,491.082,492.163,493.219,494.25,495.255,496.234,497.188,498.117,499.021,499.899,500.752,501.581,502.384,503.162,503.915,504.644,505.347,506.026,506.68,507.309,507.914,508.494,509.048,509.579,510.084,510.565,511.021,511.453,511.859,512.242,512.6,512.933,513.242,513.526,513.786,514.022,514.233,514.42,514.583,514.722,514.837,514.929,514.996,515.039,515.059
};

// Array to store lateral acceleration data (example values, replace with actual data)
double lateral_acceleration_data[data_length] = {
    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0.554,0.55,0.455,0.108,0.647,1.789,3.287,4.9,6.588,8.013,8.552,8.701,8.585,8.273,7.729,6.812,5.181,3.498,2.869,2.681,2.679,2.901,3.444,4.541,6.02,6.415,6.432,6.226,5.789,4.927,3.481,2.761,2.655,2.87,3.429,4.59,5.182,5.143,4.791,3.83,2.52,2.062,2.201,2.764,3.907,4.88,5.037,4.354,3.167,1.45,0.393,0.137,0.303,0.884,1.353,1.229,0.819,0.082,1.156,1.444,1.533,1.651,1.792,1.96,1.978,1.747,1.169,0.248,0.29,0.309,0.025,0.625,1.695,2.034,1.976,1.691,0.962,0.02,0.412,0.41,0.165,0.358,1.268,1.813,1.853,1.608,1.026,0.189,0.199,0.258,0.056,0.508,1.317,1.623,1.494,1.064,0.255,0.018,0.117,0.593,1.056,1.272,1.091,0.55,0.038,0.213,0.016,0.597,1.151,1.23,1.039,0.649,0.152,0.117,0.029,0.463,0.825,1.139,1.051,0.687,0.316,0.034,0.017,0.142,0.553,0.832,0.907,0.709,0.365,0.141,0.14,0.251,0.535,0.678,0.646,0.417,0.355,0.191,0.212,0.414,0.47,0.427,0.345,0.325,0.224,0.24,0.354,0.418,0.468,0.452,0.49,0.369,0.232,0.162,0.078,0.118,0.235,0.407,0.475,0.475,0.394,0.269,0.182,0.209,0.173,0.109,0.199,0.257,0.293,0.318,0.379,0.393,0.298,0.144,0.102,0.12,0.136,0.159,0.18,0.234,0.21,0.257,0.272,0.288,0.222,0.257,0.204,0.165,0.162,0.174,0.162,0.126,0.108,0.111,0.118,0.103,0.175,0.198,0.154,0.184,0.198,0.208,0.285,0.301,0.173,0.128,0.072,0.044,0.069,0.083,0.063,0.055,0.074,0.145,0.2,0.19,0.213,0.231,0.265,0.266,0.243,0.243,0.22,0.192,0.178,0.135,0.122,0.109,0.117,0.12,0.11,0.094,0.081,0.089,0.101,0.106,0.12,0.128,0.136,0.146,0.17,0.175,0.2,0.227,0.24,0.242,0.256,0.263,0.265,0.276,0.295
};

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 Z: ");
    Serial.println(accel_bias_z);
    Serial.print("Accel Bias X: ");
    Serial.println(accel_bias_x);

    // 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_z;
    double accelX = mpu.getAccX() - accel_bias_x;

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

    double measured_altitude = simulated_altitude;
    double measured_accelZ = accelZ;
    double measured_accelX = accelX;

    kalmanUpdate(measured_altitude, measured_accelZ, measured_accelX);

    if (!apogee_detected && velocity_estimate <= 0) {
        apogee_detected = true;
        Serial.print("Apogee detected at time: ");
        Serial.print(data_index * time_step);
        Serial.print(" s, Altitude: ");
        Serial.print(altitude_estimate);
        Serial.println(" m");
    }

    Serial.print("Simulated Altitude: ");
    Serial.print(simulated_altitude);
    Serial.print(", Estimated Altitude: ");
    Serial.print(altitude_estimate);
    Serial.print(", Estimated Lateral Distance: ");
    Serial.println(lateral_distance_estimate);

    delay(100);
}

void calibrateMPU6050() {
    long accelZ_sum = 0;
    long accelX_sum = 0;
    int samples = 1000;
    for (int i = 0; i < samples; i++) {
        mpu.update();
        accelZ_sum += mpu.getAccZ();
        accelX_sum += mpu.getAccX();
        delay(2);
    }
    accel_bias_z = accelZ_sum / samples;
    accel_bias_x = accelX_sum / samples;
}

void kalmanUpdate(double measured_altitude, double measured_accelZ, double measured_accelX) {
    // Predict step
    double dt = time_step;
    altitude_estimate += velocity_estimate * dt + 0.5 * measured_accelZ * dt * dt;
    velocity_estimate += measured_accelZ * dt;
    lateral_distance_estimate += lateral_velocity_estimate * dt + 0.5 * measured_accelX * dt * dt;
    lateral_velocity_estimate += measured_accelX * 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];
    P[2][2] += P[3][3] * dt * dt + Q[2][2];
    P[2][3] += P[3][3] * dt;
    P[3][2] += P[3][3] * dt;
    P[3][3] += Q[3][3];

    // 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;
    K[2] = P[2][2] / S;
    K[3] = P[3][2] / S;

    altitude_estimate += K[0] * y;
    velocity_estimate += K[1] * y;
    lateral_distance_estimate += K[2] * y;
    lateral_velocity_estimate += K[3] * y;

    double P00_temp = P[0][0];
    double P01_temp = P[0][1];
    double P22_temp = P[2][2];
    double P23_temp = P[2][3];

    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;
    P[2][2] -= K[2] * P22_temp;
    P[2][3] -= K[2] * P23_temp;
    P[3][2] -= K[3] * P22_temp;
    P[3][3] -= K[3] * P23_temp;
}