#include <EEPROM.h>
#include "I2Cdev.h"
#include "MPU6050_6Axis_MotionApps20.h"
#if I2CDEV_IMPLEMENTATION == I2CDEV_ARDUINO_WIRE
  #include "Wire.h"
#endif
#include <avr/wdt.h>

// Amount of readings used to average, make it higher to get more precision 
// but sketch will be slower  (default:1000)
#define CALIB_BUFFER_SIZE 1000
// Acelerometer error allowed, make it lower to get more precision, 
// but sketch may not converge  (default:8)
#define CALIB_ACCEL_DEADZONE 8
// Giro error allowed, make it lower to get 
// more precision, but sketch may not converge  (default:1)
#define CALIB_GYRO_DEADZONE 1

#define INTERRUPT_PIN 2
#define LED_PIN 13
#define CURRENT_VERSION "v1.0.0"
#define CONFIG_VERSION 0
#define SERIAL_TIMEOUT 100
#define DMP_WARMUP_TIME 10000

#define STATE_IDLE 0
#define STATE_CALIBRATION 1
#define STATE_SET_MOTORS 2
#define STATE_STOP_MOTORS 3

#define VER_MTR_PWR_PIN 4
#define HOR_MTR_PWR_PIN 5
#define VER_MTR_DIR_PIN 6
#define HOR_MTR_DIR_PIN 7

#define MTR_CTRL_VERPWRBIT 1
#define MTR_CTRL_HORPWRBIT 2
#define MTR_CTRL_VERDIRBIT 4
#define MTR_CTRL_HORDIRBIT 8

#define CMD_CALIBRATE  'c'
#define CMD_STOP_MOTOR 's'
#define CMD_SET_MOTOR  'w'

const char magicWord[5] = "DISH";
const char currVersion[11] = CURRENT_VERSION;

struct Config {
  char magic[4];
  char version[10];
  char configVersion;
  char firstSetup;
  int16_t gyroOffsetX;
  int16_t gyroOffsetY;
  int16_t gyroOffsetZ;
  int16_t accelOffsetX;
  int16_t accelOffsetY;
  int16_t accelOffsetZ;
};


MPU6050 mpu;
float ypr[3];
int mpuStatus;
int dmpStatus;
int dmpPacketSize;
uint16_t fifoCount;
uint8_t fifoBuffer[64];
float temperature;
Config trackConfig;
Quaternion q;
VectorFloat gravity;

long lastSerialTime;
bool readingSerial;

#define resetDevice() wdt_enable(WDTO_30MS); while(1) {}

void measureSensorsMean();
void saveConfig();

/*
 * MPU Interrupt
 */

volatile bool mpuInterrupt = false;     // indicates whether MPU interrupt pin has gone high
void dmpDataReady() {
  mpuInterrupt = true;
}

/*
 *  Calibration
 */

int16_t ax, ay, az,gx, gy, gz;

int meanAccelX,meanAccelY,meanAccelZ,meanGyroX,meanGyroY,meanGyroZ,state=0;
int accelOffsetX,accelOffsetY,accelOffsetZ,gyroOffsetX,gyroOffsetY,gyroOffsetZ;

void calibrate() {
  Serial.println("|0|Reseting offsets.;");
  mpu.setXGyroOffset(0);
  mpu.setYGyroOffset(0);
  mpu.setZGyroOffset(0);
  mpu.setXAccelOffset(0);
  mpu.setYAccelOffset(0);
  mpu.setZAccelOffset(0);
  
  Serial.println("|0|Reading sensor means.;");
  measureSensorsMean();

  Serial.println("|0|Calculating offsets.;");
  calculateOffsets();

  Serial.println("|0|New offsets calculated.;");
  trackConfig.gyroOffsetX = gyroOffsetX;
  trackConfig.gyroOffsetY = gyroOffsetY;
  trackConfig.gyroOffsetZ = gyroOffsetZ;

  trackConfig.accelOffsetX = accelOffsetX;
  trackConfig.accelOffsetY = accelOffsetY;
  trackConfig.accelOffsetZ = accelOffsetZ;
  
  trackConfig.firstSetup = false;

  Serial.println("|0|Saving configurations.;");
  saveConfig();
  Serial.println("|0|Resetting device.;");
  delay(100);
  resetDevice();
}

void calculateOffsets() {
  accelOffsetX=-meanAccelX/8;
  accelOffsetY=-meanAccelY/8;
  accelOffsetZ=(16384-meanAccelZ)/8;

  gyroOffsetX=-meanGyroX/4;
  gyroOffsetY=-meanGyroY/4;
  gyroOffsetZ=-meanGyroZ/4;

  while (1) { // TODO: Limit iterations
    int ready=0;
    mpu.setXAccelOffset(accelOffsetX);
    mpu.setYAccelOffset(accelOffsetY);
    mpu.setZAccelOffset(accelOffsetZ);

    mpu.setXGyroOffset(gyroOffsetX);
    mpu.setYGyroOffset(gyroOffsetY);
    mpu.setZGyroOffset(gyroOffsetZ);

    measureSensorsMean();

    if (abs(meanAccelX)<=CALIB_ACCEL_DEADZONE) {
      ready++;
    } else {
      accelOffsetX=accelOffsetX-meanAccelX/CALIB_ACCEL_DEADZONE;
    }

    if (abs(meanAccelY)<=CALIB_ACCEL_DEADZONE) {
      ready++;
    } else {
      accelOffsetY=accelOffsetY-meanAccelY/CALIB_ACCEL_DEADZONE;
    }

    if (abs(16384-meanAccelZ)<=CALIB_ACCEL_DEADZONE) {
      ready++;
    } else {
      accelOffsetZ=accelOffsetZ+(16384-meanAccelZ)/CALIB_ACCEL_DEADZONE;
    }

    if (abs(meanGyroX)<=CALIB_GYRO_DEADZONE) {
      ready++;
    } else {
      gyroOffsetX=gyroOffsetX-meanGyroX/(CALIB_GYRO_DEADZONE+1);
    }

    if (abs(meanGyroY)<=CALIB_GYRO_DEADZONE) {
      ready++;
    } else {
      gyroOffsetY=gyroOffsetY-meanGyroY/(CALIB_GYRO_DEADZONE+1);
    }

    if (abs(meanGyroZ)<=CALIB_GYRO_DEADZONE) {
      ready++;
    } else {
      gyroOffsetZ=gyroOffsetZ-meanGyroZ/(CALIB_GYRO_DEADZONE+1);
    }

    if (ready==6) break;
  }
}

void measureSensorsMean() {
  long i=0,buff_ax=0,buff_ay=0,buff_az=0,buff_gx=0,buff_gy=0,buff_gz=0;

  while (i<(CALIB_BUFFER_SIZE+101)){
    // read raw accel/gyro measurements from device
    mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);
    
    if (i>100 && i<=(CALIB_BUFFER_SIZE+100)){ //First 100 measures are discarded
      buff_ax=buff_ax+ax;
      buff_ay=buff_ay+ay;
      buff_az=buff_az+az;
      buff_gx=buff_gx+gx;
      buff_gy=buff_gy+gy;
      buff_gz=buff_gz+gz;
    }
    if (i==(CALIB_BUFFER_SIZE+100)){
      meanAccelX=buff_ax/CALIB_BUFFER_SIZE;
      meanAccelY=buff_ay/CALIB_BUFFER_SIZE;
      meanAccelZ=buff_az/CALIB_BUFFER_SIZE;
      meanGyroX=buff_gx/CALIB_BUFFER_SIZE;
      meanGyroY=buff_gy/CALIB_BUFFER_SIZE;
      meanGyroZ=buff_gz/CALIB_BUFFER_SIZE;
    }
    i++;
    delay(2); //Needed so we don't get repeated measures
  }
}


/*
 * EEPROM Manipulation
 */

void clearEEPROM() {
  for (int index = 0 ; index < EEPROM.length() ; index++) {
  EEPROM[index] = 0;
  }
}

void saveConfig() {
  EEPROM.put(0, trackConfig);
}

void loadConfig() {
  EEPROM.get(0, trackConfig);
  bool magicOk = true;
  for (int i=0; i<4; i++) {
    if (magicWord[i] != trackConfig.magic[i]) {
      magicOk = false;
    }
  }

  if (!magicOk) {
    Serial.println("|0|First setup. Cleaning EEPROM.;");
    clearEEPROM();

    for (int i=0; i<4; i++) {
      trackConfig.magic[i] = magicWord[i];
    }

    for (int i=0; i<10; i++) {
      trackConfig.version[i] = currVersion[i];
    }

    trackConfig.configVersion = CONFIG_VERSION;
    trackConfig.gyroOffsetX = 0;
    trackConfig.gyroOffsetY = 0;
    trackConfig.gyroOffsetZ = 0;
    trackConfig.accelOffsetX = 0;
    trackConfig.accelOffsetY = 0;
    trackConfig.accelOffsetZ = 0;
    trackConfig.firstSetup = true;
    saveConfig();
  }
}

/*
 * Main Code
 */

void setup() {
  #if I2CDEV_IMPLEMENTATION == I2CDEV_ARDUINO_WIRE
    Wire.begin();
    TWBR = 24; // 400kHz I2C clock (200kHz if CPU is 8MHz)
  #elif I2CDEV_IMPLEMENTATION == I2CDEV_BUILTIN_FASTWIRE
    Fastwire::setup(400, true);
  #endif

  MCUSR = 0;

  Serial.begin(115200);
  Serial.println("|0|Initializing devices.;");
  mpu.initialize();
  
  Serial.print("|1|MPUStatus|");
  Serial.print(mpu.testConnection());
  Serial.println(";");

  Serial.println("|0|Initializing DMP.;");
  dmpStatus = mpu.dmpInitialize();

  Serial.println("|0|Loading configuration.;");
  loadConfig();

  if (dmpStatus == 0) {
    dmpStatus = 1;
    Serial.println("|0|DMP Initialized.;");
  } else {
    dmpStatus = 0;
    Serial.println("|0|Failed to initialize DMP.;");
  }

  if (trackConfig.firstSetup) {
    Serial.println("|1|CALIBRATION|Requires Calibration;");
  } else {
    Serial.println("|0|Loading Accel / Gyro Calibration;");
    mpu.setXGyroOffset(trackConfig.gyroOffsetX);
    mpu.setYGyroOffset(trackConfig.gyroOffsetY);
    mpu.setZGyroOffset(trackConfig.gyroOffsetZ);
    mpu.setXAccelOffset(trackConfig.accelOffsetX);
    mpu.setYAccelOffset(trackConfig.accelOffsetY);
    mpu.setZAccelOffset(trackConfig.accelOffsetZ);
    
    mpu.setDMPEnabled(true);

    attachInterrupt(0, dmpDataReady, RISING);
    dmpPacketSize = mpu.dmpGetFIFOPacketSize();
    Serial.println("|0|Waiting for DMP warmup.;");
    lastSerialTime = millis();
    while (lastSerialTime + DMP_WARMUP_TIME > millis()); // Wait for warmpup
    Serial.println("|0|DMP Ready.;");
  }

  lastSerialTime = 0;
  readingSerial = false;
}

int currentState = STATE_IDLE;

void machineStateProcess() {
  char d;
  switch (currentState) {
    case STATE_IDLE: // Do Nothing
    break;
    case STATE_CALIBRATION:
      Serial.println("|0|Entering Calibration State. This can take some time;");
      calibrate();
      currentState = STATE_IDLE;
    break;
    case STATE_SET_MOTORS:
      d = Serial.read();
      
      Serial.print("|0|Setting motors to ");
      Serial.print(d & MTR_CTRL_VERDIRBIT > 0 ? "UP " : "DOWN ");
      Serial.print(d & MTR_CTRL_VERPWRBIT > 0 ? "ON " : "OFF ");
      Serial.print(d & MTR_CTRL_HORDIRBIT > 0 ? "LEFT " : "RIGHT ");
      Serial.println(d & MTR_CTRL_HORPWRBIT > 0 ? "ON;" : "OFF;");

      digitalWrite(VER_MTR_DIR_PIN, d & MTR_CTRL_VERDIRBIT);
      digitalWrite(HOR_MTR_DIR_PIN, d & MTR_CTRL_HORDIRBIT);
      digitalWrite(VER_MTR_PWR_PIN, d & MTR_CTRL_VERPWRBIT);
      digitalWrite(HOR_MTR_PWR_PIN, d & MTR_CTRL_HORPWRBIT);
      currentState = STATE_IDLE;
    break;
    case STATE_STOP_MOTORS:
      Serial.println("|0|Stopping Motors.;");
      digitalWrite(VER_MTR_PWR_PIN, LOW);
      digitalWrite(HOR_MTR_PWR_PIN, LOW);
      currentState = STATE_IDLE;
    break;
  }
}

void loop() {
  // Process MPU
  if (dmpStatus) {
    if (mpuInterrupt || fifoCount >= dmpPacketSize) {
      mpuInterrupt = false;
      mpuStatus = mpu.getIntStatus();

      if ((mpuStatus & 0x10) || fifoCount == 1024) {
        mpu.resetFIFO();
        Serial.println("|0|MPU FIFO Overflow!;");
      } else if (mpuStatus & 0x02) {
        while (fifoCount < dmpPacketSize) {
          fifoCount = mpu.getFIFOCount();
        }

        mpu.getFIFOBytes(fifoBuffer, dmpPacketSize);
        fifoCount -= dmpPacketSize;

        mpu.dmpGetQuaternion(&q, fifoBuffer);
        mpu.dmpGetGravity(&gravity, &q);
        mpu.dmpGetYawPitchRoll(ypr, &q, &gravity);
        temperature = mpu.getTemperature() /340.00 + 36.53;

        Serial.print("|2|");
        Serial.print(ypr[0] * 180/M_PI);
        Serial.print("|");
        Serial.print(ypr[1] * 180/M_PI);
        Serial.print("|");
        Serial.print(ypr[2] * 180/M_PI);
        Serial.println(";");

        Serial.print("|3|");
        Serial.print(temperature);
        Serial.println(";");
      }
    }
  }
  // Process Serial Data
  if (Serial.available() > 0) {
    char cmd = Serial.read();
    switch (cmd) {
      case CMD_CALIBRATE:  currentState = STATE_CALIBRATION; break;
      case CMD_STOP_MOTOR: currentState = STATE_STOP_MOTORS; break;
      case CMD_SET_MOTOR:  currentState = STATE_SET_MOTORS; break;
      default: 
        Serial.print("|0|Unknown command ");
        Serial.print(cmd);
        Serial.println(";");
    }
  } else if (lastSerialTime + SERIAL_TIMEOUT < millis() && readingSerial) {
    // Timeout
    Serial.println("|0|Serial Command Timeout;");

    readingSerial = false;
  }

  machineStateProcess();
}