#include "MPU6050.h"
#include <Servo.h>
#include "math.h"
#define CHANNEL_SIZE 6
#define IBUS_BUFFSIZE 32

// Define all the variables
int16_t ax, ay, az, gx, gy, gz;
int throttle, tailServoAngle, frontLeftSpeed, frontRightSpeed, backMotorSpeed;

static uint8_t ibusIndex=0;
static uint8_t ibus[IBUS_BUFFSIZE]={0}; 
uint16_t data[CHANNEL_SIZE];

// PID constants
float kpRoll = 0.8;
float kiRoll = 0.0;
float kdRoll = 0.0;

float kpPitch = 0.8;
float kiPitch = 0.0;
float kdPitch = 0.0;

// Target angles for balance
float targetRoll = 0.0;
float targetPitch = 0.0;

// Variables for PID control
float prevErrorRoll = 0.0;
float prevErrorPitch = 0.0;
float integralRoll = 0.0;
float integralPitch = 0.0;

// Define all ESCs
Servo frontLeft;
Servo frontRight;
Servo backMiddle;

// Define tail servo
Servo tailServo;

// Define gyro and accelerometer
MPU6050 mpu;

void setup() {

  // Set ibus port baudrate
  Serial.begin(115200);

  // Init MPU
  mpu.initialize();

  // Set builtin LED
  pinMode(13, OUTPUT);
  pinMode(3, OUTPUT);
  pinMode(9, OUTPUT);
  pinMode(10, OUTPUT);
  pinMode(11, OUTPUT);
  digitalWrite(13,0);

  // Attach servo and ESC to pins
  tailServo.attach(3);              // Attach tail servo to pin 3
  frontLeft.attach(9,1000,2000);    // Attach front left motor to 9
  frontRight.attach(10,1000,2000);  // Attach front left motor to 10
  backMiddle.attach(11,1000,2000);  // Attach front left motor to 11
  // Init all 3 ESCs
  armESC();

  // Turn on LED to signal ready status
  digitalWrite(13,1);
}

void loop() {

  // Update gyro and IBUS receiver
  read_serial();
  mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);

  throttle = get_channel(2);
  tailServoAngle =  get_channel(3);
  targetRoll = map(get_channel(0), 1000, 2000, 40, -40);
  targetPitch = map(get_channel(1), 1000, 2000, -40, 40);

  // Calculate angles
  float roll = atan2(ay, az) * 180.0 / PI;
  float pitch = atan2(ax, az) * 180.0 / PI;

  // Calculate errors
  float errorRoll = targetRoll - roll;
  float errorPitch = targetPitch - pitch;

  // Calculate PID terms
  float pidRoll = kpRoll * errorRoll + kiRoll * integralRoll + kdRoll * (errorRoll - prevErrorRoll);
  float pidPitch = kpPitch * errorPitch + kiPitch * integralPitch + kdPitch * (errorPitch - prevErrorPitch);

  // Update integral terms
  integralRoll += errorRoll;
  integralPitch += errorPitch;

  // Update previous errors
  prevErrorRoll = errorRoll;
  prevErrorPitch = errorPitch;

  // Calculate motor speeds based on PID outputs
  frontLeftSpeed = constrain(throttle + pidRoll + pidPitch, 1000, 2000);
  frontRightSpeed = constrain(throttle - pidRoll + pidPitch, 1000, 2000);
  backMotorSpeed = constrain(throttle - pidRoll - pidPitch, 1000, 2000);

  if (throttle<1010) {
      frontLeftSpeed=1000;
      frontRightSpeed=1000;
      backMotorSpeed=1000;
  }
  // Send signals to motors and servo
  frontLeft.writeMicroseconds(frontLeftSpeed);
  frontRight.writeMicroseconds(frontRightSpeed);
  backMiddle.writeMicroseconds(backMotorSpeed);
  tailServo.writeMicroseconds(tailServoAngle);
  delay(1);
}


void read_serial(void) {
  uint8_t i;
  uint16_t chksum = 0xFFFF; // Initialize chksum to a proper value
  uint16_t rxsum;
  while (Serial.available()) {
    uint8_t val = Serial.read();
    if (ibusIndex == 0 && val != 0x20) {
      continue;
    }
    if (ibusIndex == 1 && val != 0x40) {
      ibusIndex = 0;
      continue;
    }
    if (ibusIndex < IBUS_BUFFSIZE) {
      ibus[ibusIndex] = val;
    }
    ibusIndex++;
    if (ibusIndex == IBUS_BUFFSIZE) {
      ibusIndex = 0;
      for (i = 0; i < 30; i++) {
        chksum -= ibus[i];
      }
      rxsum = ibus[30] + (ibus[31] << 8);
      if (chksum == rxsum) {
        data[0] = (ibus[3] << 8) + ibus[2];
        data[1] = (ibus[5] << 8) + ibus[4];
        data[2] = (ibus[7] << 8) + ibus[6];
        data[3] = (ibus[9] << 8) + ibus[8];
        data[4] = (ibus[11] << 8) + ibus[10];
        data[5] = (ibus[13] << 8) + ibus[12];
        return;
      }
    }
  }
}

uint16_t get_channel(uint8_t channel) {
    return data[channel];
}

void armESC() {
  // Arm the ESCs
  frontLeft.writeMicroseconds(1000); 
  frontRight.writeMicroseconds(1000);
  backMiddle.writeMicroseconds(1000);
  delay(1000);
  frontLeft.writeMicroseconds(2000); 
  frontRight.writeMicroseconds(2000);
  backMiddle.writeMicroseconds(2000);
  delay(500);
  frontLeft.writeMicroseconds(1000);
  frontRight.writeMicroseconds(1000);
  backMiddle.writeMicroseconds(1000);
  delay(1000);
}