#include <Wire.h>
// =====================================================
// OLED
// =====================================================

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_ADDR 0x3C

Adafruit_SSD1306 display(
  SCREEN_WIDTH,
  SCREEN_HEIGHT,
  &Wire,
  OLED_RESET
);

// =====================================================
// MPU6050
// =====================================================

Adafruit_MPU6050 mpu;

// =====================================================
// SERVOS = VIRTUAL MOTOR OUTPUTS
// =====================================================

Servo motor1;
Servo motor2;
Servo motor3;
Servo motor4;

// Motor pins
#define M1_PIN PB1
#define M2_PIN PA0
#define M3_PIN PA1
#define M4_PIN PA4

// =====================================================
// INPUTS
// =====================================================

#define ROLL_POT  A0
#define PITCH_POT A1

#define ARM_BUTTON   PC13
#define ESTOP_BUTTON PB5

// =====================================================
// OUTPUTS
// =====================================================

#define BUZZER_PIN PA8

#define ARM_LED      PA5
#define FAILSAFE_LED PA6
#define STATUS_LED   PA7

// =====================================================
// ATTITUDE
// =====================================================

float roll = 0.0;
float pitch = 0.0;
float yaw = 0.0;

float gyroX = 0.0;
float gyroY = 0.0;
float gyroZ = 0.0;

unsigned long lastTime;

// =====================================================
// PID STRUCTURE
// =====================================================

struct PID {
  float kp;
  float ki;
  float kd;

  float integral;
  float previousError;

  float outputLimit;
};

PID rollPID  = {2.0, 0.02, 0.8, 0, 0, 45};
PID pitchPID = {2.0, 0.02, 0.8, 0, 0, 45};

// =====================================================
// SYSTEM STATUS
// =====================================================

bool armed = false;
bool failsafe = false;

bool lastArmButton = HIGH;

// =====================================================
// PID FUNCTION
// =====================================================

float calculatePID(
  PID &pid,
  float target,
  float actual,
  float dt
) {

  float error = target - actual;

  pid.integral += error * dt;

  // Integral protection
  pid.integral = constrain(pid.integral, -50.0, 50.0);

  float derivative = 0;

  if (dt > 0) {
    derivative =
      (error - pid.previousError) / dt;
  }

  float output =
      pid.kp * error +
      pid.ki * pid.integral +
      pid.kd * derivative;

  pid.previousError = error;

  output = constrain(
    output,
    -pid.outputLimit,
    pid.outputLimit
  );

  return output;
}

// =====================================================
// RESET PID
// =====================================================

void resetPID() {

  rollPID.integral = 0;
  rollPID.previousError = 0;

  pitchPID.integral = 0;
  pitchPID.previousError = 0;
}

// =====================================================
// MOTOR STOP
// =====================================================

void stopMotors() {

  motor1.write(90);
  motor2.write(90);
  motor3.write(90);
  motor4.write(90);
}

// =====================================================
// MPU6050 ATTITUDE ESTIMATION
// =====================================================

void updateAttitude(float dt) {

  sensors_event_t accel;
  sensors_event_t gyro;
  sensors_event_t temp;

  mpu.getEvent(
    &accel,
    &gyro,
    &temp
  );

  // Accelerometer angles

  float accelRoll =
    atan2(
      accel.acceleration.y,
      accel.acceleration.z
    ) * 180.0 / PI;

  float accelPitch =
    atan2(
      -accel.acceleration.x,
      sqrt(
        accel.acceleration.y *
        accel.acceleration.y +
        accel.acceleration.z *
        accel.acceleration.z
      )
    ) * 180.0 / PI;

  // Gyroscope deg/sec

  gyroX =
    gyro.gyro.x * 180.0 / PI;

  gyroY =
    gyro.gyro.y * 180.0 / PI;

  gyroZ =
    gyro.gyro.z * 180.0 / PI;

  // Gyro integration

  float gyroRoll =
    roll + gyroX * dt;

  float gyroPitch =
    pitch + gyroY * dt;

  yaw += gyroZ * dt;

  // Complementary filter

  roll =
    0.98 * gyroRoll +
    0.02 * accelRoll;

  pitch =
    0.98 * gyroPitch +
    0.02 * accelPitch;

  // Keep values reasonable

  if (yaw > 180)
    yaw -= 360;

  if (yaw < -180)
    yaw += 360;
}

// =====================================================
// OLED DISPLAY
// =====================================================

void updateOLED(
  float rollOutput,
  float pitchOutput
) {

  display.clearDisplay();

  display.setTextColor(
    SSD1306_WHITE
  );

  display.setTextSize(1);

  display.setCursor(0, 0);

  display.print("FLIGHT CONTROLLER");

  display.setCursor(0, 12);
  display.print("ROLL : ");
  display.print(roll, 1);

  display.setCursor(0, 22);
  display.print("PITCH: ");
  display.print(pitch, 1);

  display.setCursor(0, 32);
  display.print("YAW  : ");
  display.print(yaw, 1);

  display.setCursor(0, 44);
  display.print("M1:");
  display.print(motor1.read());

  display.print(" M2:");
  display.print(motor2.read());

  display.setCursor(0, 54);
  display.print("M3:");
  display.print(motor3.read());

  display.print(" M4:");
  display.print(motor4.read());

  display.display();
}

// =====================================================
// FAILSAFE
// =====================================================

void triggerFailsafe() {

  failsafe = true;
  armed = false;

  stopMotors();

  digitalWrite(
    FAILSAFE_LED,
    HIGH
  );

  digitalWrite(
    ARM_LED,
    LOW
  );

  tone(
    BUZZER_PIN,
    2000
  );

  resetPID();
}

// =====================================================
// SETUP
// =====================================================

void setup() {

  // I2C
  Wire.begin();

  // Buttons

  pinMode(
    ARM_BUTTON,
    INPUT_PULLUP
  );

  pinMode(
    ESTOP_BUTTON,
    INPUT_PULLUP
  );

  // LEDs

  pinMode(
    ARM_LED,
    OUTPUT
  );

  pinMode(
    FAILSAFE_LED,
    OUTPUT
  );

  pinMode(
    STATUS_LED,
    OUTPUT
  );

  // Buzzer

  pinMode(
    BUZZER_PIN,
    OUTPUT
  );

  // Initial LED state

  digitalWrite(
    ARM_LED,
    LOW
  );

  digitalWrite(
    FAILSAFE_LED,
    LOW
  );

  digitalWrite(
    STATUS_LED,
    HIGH
  );

  // Servo initialization

  motor1.attach(M1_PIN);
  motor2.attach(M2_PIN);
  motor3.attach(M3_PIN);
  motor4.attach(M4_PIN);

  stopMotors();

  // OLED

  if (!display.begin(
        SSD1306_SWITCHCAPVCC,
        OLED_ADDR
      )) {

    while (1) {

      digitalWrite(
        FAILSAFE_LED,
        !digitalRead(FAILSAFE_LED)
      );

      delay(200);
    }
  }

  display.clearDisplay();

  display.setTextColor(
    SSD1306_WHITE
  );

  display.setTextSize(1);

  display.setCursor(0, 0);

  display.println(
    "FLIGHT CONTROLLER"
  );

  display.println();
  display.println(
    "Initializing..."
  );

  display.display();

  delay(1000);

  // MPU6050

  if (!mpu.begin()) {

    display.clearDisplay();

    display.setCursor(0, 0);

    display.println(
      "MPU6050 ERROR!"
    );

    display.display();

    triggerFailsafe();

    while (1);
  }

  // MPU configuration

  mpu.setAccelerometerRange(
    MPU6050_RANGE_4_G
  );

  mpu.setGyroRange(
    MPU6050_RANGE_500_DEG
  );

  mpu.setFilterBandwidth(
    MPU6050_BAND_21_HZ
  );

  lastTime = micros();

  display.clearDisplay();

  display.setCursor(0, 0);

  display.println(
    "SYSTEM READY"
  );

  display.println();
  display.println(
    "DISARMED"
  );

  display.display();

  delay(1000);
}

// =====================================================
// MAIN LOOP
// =====================================================

void loop() {

  // ===================================================
  // TIME
  // ===================================================

  unsigned long now =
    micros();

  float dt =
    (now - lastTime) /
    1000000.0;

  lastTime = now;

  // Protect dt

  if (dt <= 0 || dt > 0.1)
    dt = 0.01;

  // ===================================================
  // EMERGENCY STOP
  // ===================================================

  if (
    digitalRead(
      ESTOP_BUTTON
    ) == LOW
  ) {

    triggerFailsafe();

    delay(50);

    return;
  }

  // ===================================================
  // ARM / DISARM BUTTON
  // ===================================================

  bool armButton =
    digitalRead(
      ARM_BUTTON
    );

  if (
    lastArmButton == HIGH &&
    armButton == LOW
  ) {

    if (!failsafe) {

      armed = !armed;

      resetPID();

      if (armed) {

        digitalWrite(
          ARM_LED,
          HIGH
        );

        tone(
          BUZZER_PIN,
          1000,
          100
        );

      } else {

        digitalWrite(
          ARM_LED,
          LOW
        );

        stopMotors();
      }
    }

    delay(150);
  }

  lastArmButton =
    armButton;

  // ===================================================
  // ATTITUDE ESTIMATION
  // ===================================================

  updateAttitude(dt);

  // ===================================================
  // AUTOMATIC TILT FAILSAFE
  // ===================================================

  if (
    fabs(roll) > 60 ||
    fabs(pitch) > 60
  ) {

    triggerFailsafe();

    return;
  }

  // ===================================================
  // READ PILOT SETPOINT
  // ===================================================

  int rollADC =
    analogRead(ROLL_POT);

  int pitchADC =
    analogRead(PITCH_POT);

  float targetRoll =
    map(
      rollADC,
      0,
      4095,
      -20,
      20
    );

  float targetPitch =
    map(
      pitchADC,
      0,
      4095,
      -20,
      20
    );

  // ===================================================
  // PID
  // ===================================================

  float rollCorrection = 0;

  float pitchCorrection = 0;

  if (armed) {

    rollCorrection =
      calculatePID(
        rollPID,
        targetRoll,
        roll,
        dt
      );

    pitchCorrection =
      calculatePID(
        pitchPID,
        targetPitch,
        pitch,
        dt
      );
  }

  // ===================================================
  // MOTOR MIXING
  // ===================================================

  // 90 = neutral
  // corrections are applied around neutral

  int M1 =
    90 +
    pitchCorrection +
    rollCorrection;

  int M2 =
    90 +
    pitchCorrection -
    rollCorrection;

  int M3 =
    90 -
    pitchCorrection -
    rollCorrection;

  int M4 =
    90 -
    pitchCorrection +
    rollCorrection;

  // Limit servo range

  M1 = constrain(
    M1,
    45,
    135
  );

  M2 = constrain(
    M2,
    45,
    135
  );

  M3 = constrain(
    M3,
    45,
    135
  );

  M4 = constrain(
    M4,
    45,
    135
  );

  // ===================================================
  // MOTOR OUTPUT
  // ===================================================

  if (armed && !failsafe) {

    motor1.write(M1);
    motor2.write(M2);
    motor3.write(M3);
    motor4.write(M4);

  } else {

    stopMotors();
  }

  // ===================================================
  // OLED
  // ===================================================

  updateOLED(
    rollCorrection,
    pitchCorrection
  );

  // ===================================================
  // STATUS LED
  // ===================================================

  if (!failsafe) {

    digitalWrite(
      STATUS_LED,
      HIGH
    );
  }

  delay(10);
}