#include <Wire.h>

long G[3], A[3], C[6];
unsigned long loop_timer;
// bool calibration_finished = false;
bool init_angles = false;

const int PROBES = 500;
float pitch = 0.0, roll = 0.0, accel_vector = 0.0;
float pitch_cor = 0.0, roll_cor = 0.0;
float pitch_comp = 0.0, roll_comp = 0.0;

#define GYRO_TO_DEGRESS (0.07) // mdps/LSB
#define GYRO_RESOLUTION (1 / (250 / GYRO_TO_DEGRESS))
#define RESOLUTION_TO_DEGRESS (GYRO_RESOLUTION * (PI / 180))
#define INVERT_TO_DEGRESS (1 / (PI / 180))


const float CORRECTION_RESOLUTION = 0.9996,
            COMPLIMENTARY_RESOLUTION = 0.9;

void WriteReg(uint8_t reg, uint8_t value)
{
  // HIGH_ADDR
  Wire.beginTransmission(0b1101011);
  Wire.write(reg);
  Wire.write(value);
  Wire.endTransmission();
}

void setup()
{
  Serial.begin(9600);

  Serial.println("IMU TEST");

  Wire.begin();
  TWBR = 12;

  for(int i = 0; i < 3; i++)
    G[i] = A[i] = C[i] = C[i + 3] = 0;

  WriteReg(0x10, 0b10000000); // 52hz, 8g accel
  WriteReg(0x11, 0b10001100); // 1.66kHz, 2000 dps gyro
  WriteReg(0x12, 0b00000100); // increment registery addr

  for(int i = 0; i < PROBES; i++)
  {

    if(i % 100 == 0) Serial.print(".");

    Read();

    for(int l = 0; l < 3; l++)
    {
      C[l] += G[l];
     // C[l + 3] += A[l];
    }

    delay(3);
  }
  for(int l = 0; l < 6; l++)
    C[l] /= PROBES;

  Serial.println("Calibration finished!");

  Serial.println("GYRO:");
  Serial.print("X - "); Serial.println(C[0]);
  Serial.print("Y - "); Serial.println(C[1]);
  Serial.print("Z - "); Serial.println(C[2]);

  Serial.println("ACCEL:");
  Serial.print("X - "); Serial.println(C[3]);
  Serial.print("Y - "); Serial.println(C[4]);
  Serial.print("Z - "); Serial.println(C[5]);

  // calibration_finished = true;
  loop_timer = micros();
}

void Read()
{
  Wire.beginTransmission(0b1101011);
  Wire.write(0x28);
  Wire.endTransmission();
  Wire.requestFrom((uint8_t)0b1101011, (uint8_t)6);

  while(Wire.available() < 6); // wait for readings

  for(int i = 0; i < 3; i++)
  {
    A[i] = (long)(Wire.read() << 8 | Wire.read());

   // if(calibration_finished) A[i] -= C[i];
  }

  accel_vector = sqrt(A[0] * A[0] + A[1] * A[1] + A[2] * A[2]);

  Wire.beginTransmission(0b1101011);
  Wire.write(0x22);
  Wire.endTransmission();
  Wire.requestFrom((uint8_t)0b1101011, (uint8_t)6);

  while(Wire.available() < 6); // wait for readings

  for(int i = 0; i < 3; i++)
  {
    G[i] = (long)(Wire.read() << 8 | Wire.read());

   // if(calibration_finished) G[i] -= C[i + 3];
  }
}

void loop()
{
  Read();

  pitch += (G[0] - C[0]) * GYRO_RESOLUTION;
  roll += (G[1] - C[1]) * GYRO_RESOLUTION;

  pitch -= roll * sin((G[2] - C[2]) * RESOLUTION_TO_DEGRESS);
  roll += pitch * sin((G[2] - C[2]) * RESOLUTION_TO_DEGRESS);

  if(abs(A[1]) < accel_vector)
    pitch_cor = asin(A[1] / accel_vector) * INVERT_TO_DEGRESS;

  if(abs(A[0]) < accel_vector)
    roll_cor = asin(A[0] / accel_vector) * -INVERT_TO_DEGRESS;

  pitch_cor -= C[4]; // Calibration value for A[pitch]
  roll_cor -= C[3]; // Calibration value for A[roll]

  if(init_angles)
  {
    pitch = pitch * CORRECTION_RESOLUTION + pitch_cor * (1.0 - CORRECTION_RESOLUTION);
    roll = roll * CORRECTION_RESOLUTION + roll_cor * (1.0 - CORRECTION_RESOLUTION);
  }
  else
  {
    pitch = pitch_cor;
    C[4] = pitch;
    roll = roll_cor;
    C[3] = roll;
    init_angles = true;
  }

  pitch_comp = pitch_comp * COMPLIMENTARY_RESOLUTION + pitch * (1.0 - COMPLIMENTARY_RESOLUTION);
  roll_comp = roll_comp * COMPLIMENTARY_RESOLUTION + roll * (1.0 - COMPLIMENTARY_RESOLUTION);

  Serial.println();
  Serial.print("Pitch: "); Serial.println(pitch_comp);
  Serial.print("Roll: "); Serial.println(roll_comp);
  Serial.println();

  // 250 Hz
  while(micros() - loop_timer < 4000);
  loop_timer = micros();
}