#include <SPI.h>
#include <SD.h>
#include <Servo.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>
#include <Wire.h>

const int BUZZER_PIN = 2;
const int SERVO_A_PIN = 7;
const int CHIP_SELECT = 4;

// Create objects for modules
File dataFile;
Servo A;
Adafruit_MPU6050 mpu;

// Flight stages
enum FlightStage {
  READY,
  ARM,
  ASCENT,
  DESCENT,
  LANDED
};

FlightStage current_stage = READY;

void setup() {
  Serial.begin(115200);
  delay(1000);

  pinMode(BUZZER_PIN, OUTPUT);
  digitalWrite(BUZZER_PIN, LOW);

  A.attach(SERVO_A_PIN);
  A.write(80);
  current_stage = READY;

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

  mpu.setAccelerometerRange(MPU6050_RANGE_16_G);
  mpu.setGyroRange(MPU6050_RANGE_250_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);

  delay(1000);
}

void loop() {
  switch (current_stage) {
    case READY:
      ready_stage();
      break;
    case ARM:
      arm_stage();
      break;
    case ASCENT:
      ascent_stage();
      break;
    case DESCENT:
      descent_stage();
      break;
    case LANDED:
      landed_stage();
      break;
  }
}

void beep() {
  digitalWrite(BUZZER_PIN, HIGH);
  delay(1000);
  digitalWrite(BUZZER_PIN, LOW);
}

void ready_stage() {
  Serial.println("# READY stage");
  analogWrite(LED_BUILTIN, HIGH);
  beep();

  Serial.println("# MPU6050 initialized");
  current_stage = ARM;
}

void arm_stage() {
  Serial.println("# ARM stage");
  beep();
  beep();

  int counter = 0;
  bool toggle = false;
  sensors_event_t a, g, temp;

  while (true) {
    mpu.getEvent(&a, &g, &temp);

    float x = a.acceleration.x;
    float y = a.acceleration.y;
    float z = a.acceleration.z;

    Serial.print(x); Serial.print(", ");
    Serial.print(y); Serial.print(", ");
    Serial.println(z);

    if (x > 1.0) {
      Serial.println("# Launch Detect");
      counter++;
    } else {
      counter = 0;
    }

    if (counter >= 15) {
      Serial.println("# Launching");
      break;
    }

    digitalWrite(LED_BUILTIN, toggle ? HIGH : LOW);
    toggle = !toggle;
    delay(10);
  }

  current_stage = ASCENT;
}

void ascent_stage() {
  Serial.println("# ASCENT stage");
  unsigned long StartTime = millis();

  while (true) {
    unsigned long CurrentTime = millis();
    unsigned long ElapsedTime = CurrentTime - StartTime;
    Serial.println(ElapsedTime);

    sensors_event_t a, g, temp;
    mpu.getEvent(&a, &g, &temp);

    // Add orientation failure logic here if needed

    if (ElapsedTime > 9000) {
      Serial.println("# Apogee detected by time");
      break;
    }

    delay(10);
  }

  current_stage = DESCENT;
}

void descent_stage() {
  Serial.println("# DESCENT stage");
  A.write(80);
  delay(1000);
  A.write(180);

  Serial.println("# Parachute deployed");

  unsigned long StartTime = millis();
  while (true) {
    unsigned long CurrentTime = millis();
    unsigned long ElapsedTime = CurrentTime - StartTime;

    if (ElapsedTime > 300000) {
      Serial.println("# Landing detected");
      break;
    }

    delay(1000);
  }

  current_stage = LANDED;
}

void landed_stage() {
  while (true) {
    Serial.println("# LANDED stage");
    beep();
    delay(200);
  }
}