#include <Stepper.h>
#include <Wire.h>

enum MotorPin {
  MOTOR_TR = 4,
  MOTOR_BR = 5,
  MOTOR_TL = 2,
  MOTOR_BL = 3,
};

void stepMotors() {
  digitalWrite(MOTOR_TR, HIGH);
  digitalWrite(MOTOR_BR, HIGH);
  digitalWrite(MOTOR_TL, HIGH);
  digitalWrite(MOTOR_BL, HIGH);

  digitalWrite(MOTOR_TR, LOW);
  digitalWrite(MOTOR_BR, LOW);
  digitalWrite(MOTOR_TL, LOW);
  digitalWrite(MOTOR_BL, LOW);
}

void stepMotor(MotorPin motor){
  digitalWrite(motor, HIGH);
  digitalWrite(motor, LOW);
}

void setupSensor(int addr){
  // Activate the MPU-6050
  Wire.beginTransmission(addr);
  Wire.write(0x6B);
  Wire.write(0x00);
  Wire.endTransmission();
  
  // Configure the accelerometer (+/-8g)
  Wire.beginTransmission(addr);
  Wire.write(0x1C);
  Wire.write(0x10);
  Wire.endTransmission();
  
  // Configure the gyro (500dps full scale)
  Wire.beginTransmission(addr);
  Wire.write(0x1B);
  Wire.write(0x08);
  Wire.endTransmission();
}

struct Vec2i {
  int x, y, z;
};

struct Vec2l {
  long x, y, z;
};

struct SensorData {
  Vec2l acc;
  Vec2l gyro;
  int temp;
};

// addr: 0x68
struct SensorData readSensorData(int addr){
  SensorData result = {0};
  Wire.beginTransmission(addr);
  Wire.write(0x3B);
  Wire.endTransmission();
  Wire.requestFrom(addr, 14);
  while(Wire.available() < 14);
  
  result.acc.x = Wire.read() << 8 | Wire.read();
  result.acc.y = Wire.read() << 8 | Wire.read();
  result.acc.z = Wire.read() << 8 | Wire.read();
  
  result.temp = Wire.read() << 8 | Wire.read();
  
  result.gyro.x = Wire.read() << 8 | Wire.read();
  result.gyro.y = Wire.read() << 8 | Wire.read();
  result.gyro.z = Wire.read() << 8 | Wire.read();
  
  return result;
}

#define CLAMP(mn, value, mx) ((mn) > (value) ? (mn) : (value) > (mx) ? (mx) : (value))

#define P_CONST 0.7
#define I_CONST 0.002
#define D_CONST 6.75
#define PID_MAX 250

double PID_calc(double currentAngle, double targetAngle, double* lastError, double* I){
  double error = currentAngle - targetAngle;
  *I += error * I_CONST;
  *I = CLAMP(-PID_MAX, *I, PID_MAX);

  double result = error * P_CONST + *I + (error - *lastError) * D_CONST;
  result = CLAMP(-PID_MAX, result, PID_MAX);

  *lastError = error;
  return result;
}

#define MPU6050_ADDR 0x68

Vec2l gyroError = {0};
long loopTime = 0;
void setup() {
  Serial.begin(9600);
  Wire.begin();
  pinMode(MOTOR_TR, OUTPUT);
  pinMode(MOTOR_BR, OUTPUT);
  pinMode(MOTOR_TL, OUTPUT);
  pinMode(MOTOR_BL, OUTPUT);

  setupSensor(MPU6050_ADDR);

#define CALIBRATION_INTERATION_COUNT 1 // TODO: change back to 2000
  for(int i = 0; i < CALIBRATION_INTERATION_COUNT; i++){
    if(i % 15 == 0) digitalWrite(LED_BUILTIN, HIGH);
    SensorData data = readSensorData(MPU6050_ADDR);
    gyroError.x += data.gyro.x;
    gyroError.y += data.gyro.y;
    gyroError.z += data.gyro.z;
    digitalWrite(LED_BUILTIN, LOW);
    delay(3);
  }
  gyroError.x /= CALIBRATION_INTERATION_COUNT;
  gyroError.y /= CALIBRATION_INTERATION_COUNT;
  gyroError.z /= CALIBRATION_INTERATION_COUNT;

  loopTime = micros();
}

void debugPrint(struct Vec2l data){
  Serial.print("(x: ");
  Serial.print(data.x);
  Serial.print(", y: ");
  Serial.print(data.y);
  Serial.print(", z: ");
  Serial.print(data.z);
  Serial.println(")");
}

bool firstFrame = true;
double anglePitch = 0;
double angleRoll  = 0;
double currentPitch = 0.0f;
double currentRoll  = 0.0f;
double lastPitchError;
double lastRollError;
double pitchAccumulatedError = 0.0f;
double rollAccumulatedError = 0.0f;
void loop() {
  SensorData data = readSensorData(MPU6050_ADDR);

  data.gyro.x -= gyroError.x;
  data.gyro.y -= gyroError.y;
  data.gyro.z -= gyroError.z;

  // Gyro angle calculations
  // 0.0000611 = 1 / (250Hz / 65.5)
  anglePitch += data.gyro.x * 0.0000611f;
  angleRoll  += data.gyro.y * 0.0000611f;

  // 0.000001066 = 0.0000611 * (3.142(PI) / 180degr)
  anglePitch += angleRoll  * sin(data.gyro.z * 0.000001066f);
  angleRoll  -= anglePitch * sin(data.gyro.z * 0.000001066f);
  
  // Accelerometer angle calculations
  long xSqrd = data.acc.x * data.acc.x;
  long ySqrd = data.acc.y * data.acc.y;
  long zSqrd = data.acc.z * data.acc.z;
  double totalVector = sqrt(xSqrd + ySqrd + zSqrd);
  
  // 57.296 = 1 / (3.142 / 180)
  double anglePitchAcc = asin((double)data.acc.y / totalVector) *  57.296f;
  double angleRollAcc  = asin((double)data.acc.x / totalVector) * -57.296f;

  // NOTE: Correcting for the error meassured on the MPU6050 when placed horizontal
  anglePitchAcc -= 0.0f;
  angleRollAcc  -= 0.0f;

  if(firstFrame){
    anglePitch = anglePitchAcc;
    angleRoll = angleRollAcc;
    firstFrame = false;
  }else{
    anglePitch = anglePitch * 0.9996f + anglePitchAcc * 0.0004f;
    angleRoll  = angleRoll  * 0.9996f + angleRollAcc  * 0.0004f;
  }

  // To dampen the pitch and roll angles a complementary filter is used
  currentPitch = currentPitch * 0.9f + anglePitch * 0.1f;
  currentRoll  = currentRoll  * 0.9f + angleRoll  * 0.1f;
  // Serial.print("currentPitch: "); Serial.println(currentPitch);
  // Serial.print("currentRoll: "); Serial.println(currentRoll);

  // auto leveling using the PID controller
  double target = 0.0f;
  double pitchResult = PID_calc(currentPitch, target, &lastPitchError, &pitchAccumulatedError);
  double rollResult  = PID_calc(currentRoll,  target, &lastRollError, &rollAccumulatedError);
  // Serial.print("pitchResult: "); Serial.println(pitchResult);
  // Serial.print("rollResult: "); Serial.println(rollResult);

  // Write the speed
  double speed = 100.0f;
  for(int i = 0; i < (int)(speed + rollResult - pitchResult); i++) stepMotor(MOTOR_TR);
  for(int i = 0; i < (int)(speed - rollResult - pitchResult); i++) stepMotor(MOTOR_TL);
  for(int i = 0; i < (int)(speed + rollResult + pitchResult); i++) stepMotor(MOTOR_BR);
  for(int i = 0; i < (int)(speed - rollResult + pitchResult); i++) stepMotor(MOTOR_BL);

  // Serial.println("--------------------");
  // Serial.print("TR: "); Serial.println(speed + rollResult - pitchResult);
  // Serial.print("TL: "); Serial.println(speed - rollResult - pitchResult);
  // Serial.print("BR: "); Serial.println(speed + rollResult + pitchResult);
  // Serial.print("BL: "); Serial.println(speed - rollResult + pitchResult);

  // delay
  while(micros() - loopTime < 4000);
  loopTime = micros();
}