#include <Arduino.h>
#include <SD.h>
#include <SPI.h>
#include <Servo.h>
#include <PID_v1.h>
#include <Wire.h>
#include <MPU6050.h>
#include <avr/wdt.h>
#define BLINK_RATE 200

//=== Available modes ===
enum RocketMode {
  MODE_IDLE,
  MODE_ALIGN_SERVO,
  MODE_GROUND_TEST,
  MODE_FLIGHT,
  MODE_FIN_TEST
};

RocketMode currentMode = MODE_IDLE;
RocketMode lastMode = (RocketMode) - 1;

//=== Objects ===
MPU6050 mpu;
Servo servo1, servo2, servo3, servo4;
File logFile;

//=== PID and variables ===
double inputX, outputX, setpointX = 0;
double inputY, outputY, setpointY = 0;
double inputZ, outputZ, setpointZ = 0;

double KpX = 4.5, KiX = 4.5, KdX = 0.45; //Kp = 4.5, Ki = 5.4, Kd = 0.45
double KpY = 4.5, KiY = 4.5, KdY = 0.45;
double KpZ = 1.0, KiZ = 0, KdZ = 0.30; //Kp = 1, Ki = 0, Kd = 0.3

PID PIDx(&inputX, &outputX, &setpointX, KpX, KiX, KdX, DIRECT);
PID PIDy(&inputY, &outputY, &setpointY, KpY, KiY, KdY, DIRECT);
PID PIDz(&inputZ, &outputZ, &setpointZ, KpZ, KiZ, KdZ, DIRECT);

float gyroX, gyroY, gyroZ;
float gyroXFiltered, gyroYFiltered, gyroZFiltered;
float accX, accY, accZ;
float angleX = 0, angleY = 0, angleZ = 0;
float offsetX = 0, offsetY = 0;
float launchAccl = 0;

int servo1Angle = 0, servo2Angle = 0, servo3Angle = 0, servo4Angle = 0;
int servo1Default = 90, servo2Default = 90, servo3Default = 90, servo4Default = 90;


//=== State ===
bool launched = false;
bool calibrating = false;
bool offsetting = false;
bool MPUnotFound = true;
bool SDnotFound = true;

// === Low-pass filter ===
float alpha = 1;

unsigned long previousTime = 0;

//=== Functions prototypes ===
void enterMode(RocketMode mode);
void updateMode(RocketMode mode);
void checkSerial();
void attachServos();
void setupMPU();
void calibrateMPU();
void calibrateInclination();
void setupPID();
float updateTime();
bool checkInterval(unsigned long intervalMs);
void checkLaunch();
void readGyroAngles(float elapsedTime);
void readAccelerometer();
void updatePID();
void alignServo();
void finTest();
void computeServoAngles();
void writeServoAngles();
void SDsetup();
void dataLogger();
void printDebug();
void blinkIfAlive();


//=== Setup ===
void setup() {
  pinMode(LED_BUILTIN, OUTPUT);
  Wire.begin();
  Serial.begin(115200);

  attachServos();
  setupMPU();
  SDsetup();
  alignServo();

  previousTime = millis();
}

//=== Main loop ===
void loop() {
  blinkIfAlive();
  checkSerial();

  if (currentMode != lastMode) {
    enterMode(currentMode);
    lastMode = currentMode;
  }

  updateMode(currentMode);
}

void blinkIfAlive() {
  static uint32_t nextBlink = 0;
  if (millis() > nextBlink) {
    digitalWrite(LED_BUILTIN, !digitalRead(LED_BUILTIN)); // pinMode(LED_BUILTIN, OUTPUT) in setup()
    nextBlink += BLINK_RATE;  // #define BLINK_RATE 200 or whatever
  }
}

//=== Modes management ===
void enterMode(RocketMode mode) {
  switch (mode) {
    case MODE_IDLE:
      calibrateMPU();
      printDebug();
      break;

    case MODE_ALIGN_SERVO:
      printDebug();
      break;

    case MODE_GROUND_TEST:
      KiX = 0; //Set correct Ki values for ground test
      KiY = 0;

      calibrateMPU();
      calibrateInclination();
      setupPID();
      launched = false;
      launchAccl = 1.5; //Set launch acceleration threshold
      printDebug();
      break;

    case MODE_FLIGHT:
      KiX = 5.4; //Set correct Ki values for flight
      KiY = 5.4;

      calibrateMPU();
      calibrateInclination();
      setupPID();
      launched = false;
      launchAccl = 1.5;
      printDebug();
      break;

    case MODE_FIN_TEST:
      printDebug();
      break;
  }
}

void updateMode(RocketMode mode) { //Fake loop for modes
  float elapsedTime = updateTime();

  switch (mode) {
    case MODE_IDLE:
      break;

    case MODE_ALIGN_SERVO:
      printDebug();
      alignServo();
      break;

    case MODE_GROUND_TEST:
      if (!launched) {
        checkLaunch();
        return;
      }

      readGyroAngles(elapsedTime);
      readAccelerometer();
      updatePID();
      computeServoAngles();
      writeServoAngles();
      printDebug();
      break;

    case MODE_FLIGHT:
      if (!launched) {
        checkLaunch();
        return;
      }

      readGyroAngles(elapsedTime);
      readAccelerometer();
      updatePID();
      computeServoAngles();
      writeServoAngles();
      dataLogger();
      break;

    case MODE_FIN_TEST:
      printDebug();
      finTest();
      break;
  }
}

void checkSerial() {
  if (Serial.available()) {
    String cmd = Serial.readStringUntil('\n');
    cmd.trim();

    //=== Virtual reset ===
    if (cmd == "RESET") {
      Serial.println(F("Riavvio..."));
      delay(100);
      wdt_enable(WDTO_15MS);
      while (1) {}
    }
    //=== Commands for changing modes ===
    if (cmd == "0") currentMode = MODE_IDLE;
    if (cmd == "1") currentMode = MODE_ALIGN_SERVO;
    if (cmd == "2") currentMode = MODE_GROUND_TEST;
    if (cmd == "3") currentMode = MODE_FLIGHT;
    if (cmd == "4") currentMode = MODE_FIN_TEST;

//MODE_ALIGN_SERVO Commands
    int val = cmd.substring(3).toInt();
    if (val != 0 || cmd.substring(3, 4) == "0") { // optionally: || cmd.substring(3, 5) == " 0"
      val = constrain(val, 0, 180);
      if (currentMode == MODE_ALIGN_SERVO) {
        if (cmd.startsWith("S1:")) {
          servo1Default = val;
        } else if (cmd.startsWith("S2:")) {
          servo2Default = val;
        } else if (cmd.startsWith("S3:")) {
          servo3Default = val;
        } else if (cmd.startsWith("S4:")) {
          servo4Default = val;
        }
      }
    }
  }
}

//=== Functions implementations ===
void attachServos() {
  //Attach servos to pins
  servo1.attach(3);
  servo2.attach(5);
  servo3.attach(6);
  servo4.attach(9);
}

void setupMPU() {
  //Setup MPU6050
  Serial.println(F("|*Avvio MPU6050*|"));
  mpu.initialize();
  if (mpu.testConnection()) {
    Serial.println(F("MPU6050 trovato!"));
    MPUnotFound = false;
    printDebug();
  } else {
    Serial.println(F("MPU6050 non trovato!"));
    MPUnotFound = true;
    printDebug();
    while (1);
  }
}

void calibrateMPU() {
  //Calibrate MPU6050 gyroscope
  calibrating = true;
  printDebug();
  Serial.println(F("Tenere il razzo fermo!"));
  mpu.CalibrateGyro();
  Serial.println(F("Calibrazione completata"));
  calibrating = false;
  printDebug();
}

void calibrateInclination() {
  //Calibrate inclination offsets
  Serial.println(F("Calibro inclinazione rampa..."));
  offsetting = true;
  printDebug();
  int samples = 100; //Number of samples for averaging
  long accXsum = 0, accYsum = 0, accZsum = 0;

  for (int i = 0; i < samples; i++) { //Sum samples
    accXsum += mpu.getAccelerationX();
    accYsum += mpu.getAccelerationY();
    accZsum += mpu.getAccelerationZ();
    delay(5);
  }

  //Calculate average
  float accX = accXsum / samples;
  float accY = accYsum / samples;
  float accZ = accZsum / samples;

  // Normalize accelerometer values
  accX /= 16384.0;
  accY /= 16384.0;
  accZ /= 16384.0;

  // Calculate offsets
  offsetX = atan2(-accX, sqrt(accY * accY + accZ * accZ)) * RAD_TO_DEG;
  offsetY = atan2(accY, accZ) * RAD_TO_DEG;

  angleX = offsetX;
  angleY = offsetY;

  Serial.print(F("Offset X: ")); Serial.println(offsetX);
  Serial.print(F("Offset Y: ")); Serial.println(offsetY);
  delay(3000);
  offsetting = false;
  printDebug();
}

void setupPID() {
  //Setup PID controllers
  PIDx.SetMode(AUTOMATIC); PIDx.SetOutputLimits(-20, 20);
  PIDy.SetMode(AUTOMATIC); PIDy.SetOutputLimits(-20, 20);
  PIDz.SetMode(AUTOMATIC); PIDz.SetOutputLimits(-20, 20);
}

float updateTime() {
  //Update elapsed time
  unsigned long currentTime = millis();
  float elapsed = (currentTime - previousTime) / 1000.0;
  previousTime = currentTime;
  return elapsed;
}

bool checkInterval(unsigned long intervalMs) {
  //Virtual configurable clock
  static unsigned long previousCheck = 0;
  unsigned long now = millis();

  if (now - previousCheck >= intervalMs) {
    previousCheck = now;
    return true;
  }
  return false;
}

void checkLaunch() {
  //Check if the rocket is launched based on accelerometer data
  readAccelerometer();
  if (accZ >= launchAccl) {
    launched = true;
    Serial.println(F(">>>>> LANCIO <<<<<"));
    printDebug();
  }
}

void readGyroAngles(float elapsedTime) {
  //Read gyro values
  gyroX = mpu.getRotationX() / 131.0;
  gyroY = mpu.getRotationY() / 131.0;
  gyroZ = mpu.getRotationZ() / 131.0;

  // === Low-pass filter ===
  gyroXFiltered = alpha * gyroX + (1 - alpha) * gyroXFiltered;
  gyroYFiltered = alpha * gyroY + (1 - alpha) * gyroYFiltered;
  gyroZFiltered = alpha * gyroZ + (1 - alpha) * gyroZFiltered;


  // === Angles calculation ===
  angleX += gyroXFiltered * elapsedTime;
  angleY += gyroYFiltered * elapsedTime;
  angleZ += gyroZFiltered * elapsedTime;
}

void readAccelerometer() {
  //Read accelerometer values
  accX = mpu.getAccelerationX() / 16384.0;
  accY = mpu.getAccelerationY() / 16384.0;
  accZ = mpu.getAccelerationZ() / 16384.0;
}

void updatePID() {
  //Update PID inputs
  inputX = angleX;
  inputY = angleY;
  inputZ = gyroZFiltered;
  PIDx.Compute();
  PIDy.Compute();
  PIDz.Compute();
}

void alignServo() {
  //Align servos to default positions
  servo1.write(servo1Default);
  servo2.write(servo2Default);
  servo3.write(servo3Default);
  servo4.write(servo4Default);
}

void finTest() {
  const int delta = 20;
  const int snapDelay = 350; // ms, time for snap movements
  const int smoothDelay = 30; // ms, time for smooth movements

  // 1. Snap singole movements
  for (int i = 0; i < 4; i++) {
    int *servoDefault[4] = {&servo1Default, &servo2Default, &servo3Default, &servo4Default};
    int *servoAngle[4] = {&servo1Angle, &servo2Angle, &servo3Angle, &servo4Angle};

    // +20°
    *servoAngle[i] = constrain(*servoDefault[i] + delta, 0, 180);
    writeServoAngles();
    printDebug();
    delay(snapDelay);

    // -20°
    *servoAngle[i] = constrain(*servoDefault[i] - delta, 0, 180);
    writeServoAngles();
    printDebug();
    delay(snapDelay);

    // Default
    *servoAngle[i] = *servoDefault[i];
    writeServoAngles();
    printDebug();
    delay(snapDelay);
  }

  // 2. Snap pair movements
  // S1/S3 opposite
  servo1Angle = constrain(servo1Default + delta, 0, 180);
  servo3Angle = constrain(servo3Default - delta, 0, 180);
  writeServoAngles();
  printDebug();
  delay(snapDelay);

  servo1Angle = servo1Default;
  servo3Angle = servo3Default;
  writeServoAngles();
  printDebug();
  delay(snapDelay);

  // S2/S4 opposite
  servo2Angle = constrain(servo2Default + delta, 0, 180);
  servo4Angle = constrain(servo4Default - delta, 0, 180);
  writeServoAngles();
  printDebug();
  delay(snapDelay);

  servo2Angle = servo2Default;
  servo4Angle = servo4Default;
  writeServoAngles();
  printDebug();
  delay(snapDelay);

  // 3. Smooth movement
  for (int angle = delta; angle >= -delta; angle -= 1) {
    servo1Angle = constrain(servo1Default + angle, 0, 180);
    servo2Angle = constrain(servo2Default + angle, 0, 180);
    servo3Angle = constrain(servo3Default + angle, 0, 180);
    servo4Angle = constrain(servo4Default + angle, 0, 180);
    writeServoAngles();
    printDebug();
    delay(smoothDelay);
  }
  for (int angle = -delta; angle <= delta; angle += 1) {
    servo1Angle = constrain(servo1Default + angle, 0, 180);
    servo2Angle = constrain(servo2Default + angle, 0, 180);
    servo3Angle = constrain(servo3Default + angle, 0, 180);
    servo4Angle = constrain(servo4Default + angle, 0, 180);
    writeServoAngles();
    printDebug();
    delay(smoothDelay);
  }

  //Back to default positions
  servo1Angle = servo1Default;
  servo2Angle = servo2Default;
  servo3Angle = servo3Default;
  servo4Angle = servo4Default;
  writeServoAngles();
  printDebug();
  delay(500);

  Serial.println(F("Fin test completato! Ritorno in idle."));
  currentMode = MODE_IDLE; //Back to idle mode
}

void computeServoAngles() {
  //Mixing matrix
  servo1Angle = servo1Default + (+1 * outputY) + (-1 * outputZ);
  servo2Angle = servo2Default + (+1 * outputX) + (+1 * outputZ);
  servo3Angle = servo3Default + (-1 * outputY) + (+1 * outputZ);
  servo4Angle = servo4Default + (-1 * outputX) + (-1 * outputZ);
}

void writeServoAngles() {
  servo1.write(servo1Angle);
  servo2.write(servo2Angle);
  servo3.write(servo3Angle);
  servo4.write(servo4Angle);
}

void SDsetup() {
  //Iniitialization
  printDebug();
  Serial.println(F("Inizializzazione SD..."));
  if (!SD.begin(10)) {
    Serial.println(F("Inizializzazione SD fallita!"));
    SDnotFound = true;
    printDebug();
    while (1);
  }
  Serial.println(F("SD inizializzata correttamente."));
  SDnotFound = false;
  printDebug();

  //Opening file and writing header
  logFile = SD.open("log.txt", FILE_WRITE);
  if (logFile) {
    logFile.println(F("Stinger logging system online"));
    logFile.println();
    logFile.println(F("Time,AX,AY,AZ,GX,GY,GZ,GXF,GYF,GZF,OUTX,OUTY,OUTZ,ANGX,ANGY,ANGZ,S1,S2,S3,S4,OFFX,OFFY,LAUNCHED,MODE"));
    logFile.flush(); //Ensure data is written to SD card
  }
}

void dataLogger() {
  if (checkInterval(50)) { //Write every 50ms

    logFile.print(millis()); logFile.print(",");
    logFile.print(accX); logFile.print(",");
    logFile.print(accY); logFile.print(",");
    logFile.print(accZ); logFile.print(",");
    logFile.print(gyroX); logFile.print(",");
    logFile.print(gyroY); logFile.print(",");
    logFile.print(gyroZ); logFile.print(",");
    logFile.print(gyroXFiltered); logFile.print(",");
    logFile.print(gyroYFiltered); logFile.print(",");
    logFile.print(gyroZFiltered); logFile.print(",");
    logFile.print(outputX); logFile.print(",");
    logFile.print(outputY); logFile.print(",");
    logFile.print(outputZ); logFile.print(",");
    logFile.print(angleX); logFile.print(",");
    logFile.print(angleY); logFile.print(",");
    logFile.print(angleZ); logFile.print(",");
    logFile.print(servo1Angle); logFile.print(",");
    logFile.print(servo2Angle); logFile.print(",");
    logFile.print(servo3Angle); logFile.print(",");
    logFile.print(servo4Angle); logFile.print(",");
    logFile.print(offsetX); logFile.print(",");
    logFile.print(offsetY); logFile.print(",");
    logFile.print(launched); logFile.print(",");
    logFile.print(currentMode); logFile.println(); //Close line
    logFile.flush(); //Ensure data is written to SD card
  }
}

void printDebug() {
  //Print data for debugging
  Serial.print(F("AX:")); Serial.print(accX, 2); Serial.write(' ');
  Serial.print(F("AY:")); Serial.print(accY, 2); Serial.write(' ');
  Serial.print(F("AZ:")); Serial.print(accZ, 2); Serial.write(' ');
  Serial.print(F("GX:")); Serial.print(gyroX, 2); Serial.write(' ');
  Serial.print(F("GY:")); Serial.print(gyroY, 2); Serial.write(' ');
  Serial.print(F("GZ:")); Serial.print(gyroZ, 2); Serial.write(' ');
  Serial.print(F("GXF")); Serial.print(gyroXFiltered, 2); Serial.write(' ');
  Serial.print(F("GYF")); Serial.print(gyroYFiltered, 2); Serial.write(' ');
  Serial.print(F("GZF")); Serial.print(gyroZFiltered, 2); Serial.write(' ');
  Serial.print(F("OUTX:")); Serial.print(outputX, 2); Serial.write(' ');
  Serial.print(F("OUTY:")); Serial.print(outputY, 2); Serial.write(' ');
  Serial.print(F("OUTZ:")); Serial.print(outputZ, 2); Serial.write(' ');
  Serial.print(F("ANGX:")); Serial.print(angleX, 2); Serial.write(' ');
  Serial.print(F("ANGY:")); Serial.print(angleY, 2); Serial.write(' ');
  Serial.print(F("ANGZ:")); Serial.print(angleZ, 2); Serial.write(' ');
  Serial.print(F("OFFSETTING:")); Serial.print(offsetting); Serial.write(' ');
  Serial.print(F("CALIBRATING:")); Serial.print(calibrating); Serial.write(' ');
  Serial.print(F("MPU:")); Serial.print(MPUnotFound); Serial.write(' ');
  Serial.print(F("SD:")); Serial.print(SDnotFound); Serial.write(' ');
  Serial.print(F("MODE:")); Serial.print(currentMode); Serial.write(' ');
  Serial.print(F("LAUNCHED:")); Serial.print(launched); Serial.write(' ');
  Serial.print(F("OFFX:")); Serial.print(offsetX, 2); Serial.write(' ');
  Serial.print(F("OFFY:")); Serial.print(offsetY, 2); Serial.write(' ');
  Serial.print(F("S1:")); Serial.print(servo1Angle); Serial.write(' ');
  Serial.print(F("S2:")); Serial.print(servo2Angle); Serial.write(' ');
  Serial.print(F("S3:")); Serial.print(servo3Angle); Serial.write(' ');
  Serial.print(F("S4:")); Serial.print(servo4Angle); Serial.write('\n');
}