#include <Wire.h>
#include <LSM303.h>
#include <L3G.h>
#include "Kalman.h"

#define LED 13

//You will need to find your own magnetometer calibration values, use the sample code that came with the LSM303 library!!!

#define MAG_X_MIN -2752
#define MAG_Y_MIN -3040
#define MAG_Z_MIN -2503
#define MAG_X_MAX +3741
#define MAG_Y_MAX +3053
#define MAG_Z_MAX +3281

#define ROLL_LIMIT 90.0f
#define PITCH_LIMIT 180.0f
#define YAW_LIMIT 180.0f

Kalman kalmanZ, kalmanX, kalmanY;

char report[32];
char inByte;
unsigned long currentTime, lastTime;
float dT;

float roll, pitch, yaw;                   // Roll and pitch calculated from accelerometer, yaw from magnetometer
float gyroXangle, gyroYangle, gyroZangle; // Angle calculated with gyroscope
float gyroXrate, gyroYrate, gyroZrate;    // Gyroscope movement rate, in degrees per second
float compAngleX, compAngleY, compAngleZ; // Calculated angle using a complementary filter, YOU WILL NEED TO DO THIS
float kalAngleX, kalAngleY, kalAngleZ;    // Calculated angle using a Kalman filter, YOU WILL NEED TO DO THIS

LSM303 compass;
L3G gyro;

void setup() {
  pinMode(LED, OUTPUT);

  Serial.begin(1000000);
  Wire.begin();
  Wire.setClock(400000L);

  compass.init();
  compass.enableDefault();
  // Set accelerometer data rate to 800Hz
  // 0b10010111 = 0x97
  compass.writeReg(0x20, 0x97);
  // Set magnetometer data rate to 100Hz
  // 0b01110100 = 0x74
  compass.writeReg(0x24, 0x74);
  compass.m_min = (LSM303::vector<int16_t>){MAG_X_MIN, MAG_Y_MIN, MAG_Z_MIN};
  compass.m_max = (LSM303::vector<int16_t>){MAG_X_MAX, MAG_Y_MAX, MAG_Z_MAX};

  if (!gyro.init())
  {
    Serial.println("Failed to autodetect gyro type!");
    while (1);
  }
  gyro.enableDefault();
  // Set gyro data rate to 800Hz
  // 0b11111111 = 0xFF
  gyro.writeReg(0x20, 0xFF);

  roll = 0;
  pitch = 0;
  yaw = 0;

  lastTime = micros();
  compass.read();
  updateRollPitch(&compass, &roll, &pitch);
  updateYaw(&compass, &yaw);

  gyroXangle = roll;
  gyroYangle = pitch;
  gyroZangle = yaw;

  compAngleX = roll;
  compAngleY = pitch;
  compAngleZ = yaw;
}

void loop() {
  currentTime = micros();
  compass.read();
  gyro.read();
  updateRollPitch(&compass, &roll, &pitch);
  updateYaw(&compass, &yaw);

  dT = (currentTime - lastTime) / 1000000.0f;
  lastTime = currentTime;

  gyroXrate = gyroToDegPerSec(gyro.g.x);
  gyroYrate = gyroToDegPerSec(gyro.g.y);
  gyroZrate = gyroToDegPerSec(gyro.g.z);

  // Fix xrate and zrate
  if (abs(pitch) > 90) gyroXrate = -gyroXrate;
  else gyroZrate = -gyroZrate;

  gyroXangle += gyroXrate * dT;
  gyroYangle += gyroYrate * dT;
  gyroZangle += gyroZrate * dT;

  kalmanX.setAngle(roll);
  kalmanY.setAngle(roll);
  kalmanZ.setAngle(roll);

  fixAngle(&gyroYangle, &pitch, PITCH_LIMIT);
  fixAngle(&gyroZangle, &yaw, YAW_LIMIT);

  // You will need to calculate these
  //compAngleX = 0.98 * (compAngleX + gyroXrate * dT) + 0.02 * roll; // Calculate the angle using a Complimentary filter
  //compAngleY = 0.98 * (compAngleY + gyroYrate * dT) + 0.02 * pitch;
  //compAngleZ = 0.98 * (compAngleZ + gyroZrate * dT) + 0.02 * yaw;

  // You will need to calculate these
  kalAngleX = kalmanX.getAngle(roll, gyroXrate, dT); // Calculate the angle using a Kalman filter
  kalAngleY = kalmanY.getAngle(pitch, gyroYrate, dT);
  kalAngleZ = kalmanZ.getAngle(yaw, gyroZrate, dT);

  if (Serial.available()) respondToSerial();
}

void respondToSerial() {
  digitalWrite(LED, !digitalRead(LED));
  inByte = Serial.read();

    switch (inByte) {
      case 'a':
        snprintf(report, sizeof(report), "<A:%d:%d:%d>",
          compass.a.x, compass.a.y, compass.a.z);
        Serial.print(report);
        break;
      case 'm':
        snprintf(report, sizeof(report), "<M:%d:%d:%d>",
          compass.m.x, compass.m.y, compass.m.z);
        Serial.print(report);
        break;
      case 'g':
        snprintf(report, sizeof(report), "<G:%d:%d:%d>",
          gyro.g.x, gyro.g.y, gyro.g.z);
        Serial.print(report);
        break;
      case 'r':
        floatsToSerial('R', roll, pitch, yaw);
        break;
      case 'd':
        floatsToSerial('D', gyroXangle, gyroYangle, gyroZangle);
        break;
      case 'c':
        floatsToSerial('C', compAngleX, compAngleY, compAngleZ);
        break;
      case 'k':
        floatsToSerial('K', kalAngleX, kalAngleY, kalAngleZ);
        break;
      case 'x':
        //RESET ANGLES
        gyroXangle = roll;
        gyroYangle = pitch;
        gyroZangle = yaw;
        compAngleX = roll;
        compAngleY = pitch;
        compAngleZ = yaw;
        break;
      default:
        Serial.print("<E:0:0:0>");
        break;
    }
}

void floatsToSerial(const char tag, float first, float second, float third) {
  String output = '<' + String(tag) + ':' + String(first, 3)+ ':' + String(second, 3)+ ':' + String(third, 3) + '>';
  Serial.print(output);
}

float gyroToDegPerSec(int16_t gyro) {
  return gyro * 0.00875f;
}

void updateRollPitch(LSM303 *compass, float *roll, float *pitch) {
  *roll = atan(compass->a.y / sqrt((float)compass->a.x * (float)compass->a.x + (float)compass->a.z * (float)compass->a.z)) * RAD_TO_DEG;
  *pitch = atan2(-compass->a.x, compass->a.z) * RAD_TO_DEG;
}

void updateYaw(LSM303 *compass, float *yaw) {
  float heading = compass->heading();
  if (heading > YAW_LIMIT) heading -= 360.0f;
  *yaw = heading;
}

void fixAngle(float *angle, float *reference, float limit) {
  if (abs(*angle - *reference) > limit * 1.5f) {
    if (*angle > *reference) *angle -= 2*limit;
    else *angle += 2*limit;
  }
}