#include <Adafruit_MPU6050.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_Sensor.h>
#include <string.h>

Adafruit_MPU6050 mpu;
Adafruit_SSD1306 display = Adafruit_SSD1306(128, 32, &Wire);
 
const int button_pin_red = 2;
const int button_pin_green = 3;
int button_state_r = 0;
int button_state_g = 0;
int button_switch = 0;

long max_acc = 0;

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

  pinMode(button_pin_red, INPUT);
  pinMode(button_pin_green, INPUT);

  delay(200);

  if (!mpu.begin()) {
    Serial.println("Sensor init failed");
    while (1)
      yield();
  }

  Serial.println("Found a MPU-6050 sensor");

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x32
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
  }

  display.display();
  delay(500); // Pause for 2 seconds
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setRotation(0);
}

void loop() {
  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);

  double acc_mag = sqrt(a.acceleration.x*a.acceleration.x + a.acceleration.y*a.acceleration.y + a.acceleration.z*a.acceleration.z);

  if(acc_mag > max_acc) {
    max_acc = acc_mag;
  }

  button_state_r = digitalRead(button_pin_red);
  button_state_g = digitalRead(button_pin_green);

  display.clearDisplay();
  display.setCursor(0, 0);
  delay(5);

  display.print("Rocket Accelerometer");
  display.println();

  if(button_state_g == 1) {
    max_acc = 0;
    delay(200);
  }

  if(button_state_r == 1) {
    button_switch++;
    delay(200);
  }
  delay(5);

  if(button_switch > 2) {
    button_switch = 0;
  }

  delay(5);

  if(button_switch == 0) {
    display_acc(a, acc_mag);
  } else if(button_switch == 1) {
    display_gyro(g);
  } else if(button_switch == 2) {
    display_max_acc(max_acc);
  }

  display.display();
  delay(100);
}

void display_acc(sensors_event_t a, double acc_mag) {
  display.println("Accelerometer - m/s^2");
  display.print("X: ");
  display.print(a.acceleration.x, 1);
  display.print(", ");
  display.print("Y: ");
  display.print(a.acceleration.y, 1);
  display.print(", ");
  display.println("");
  display.print("Z: ");
  display.print(a.acceleration.z, 1);
  display.print(", Mag: ");
  display.print(acc_mag, 1);
  display.println("");
}

void display_gyro(sensors_event_t g) {
  display.println("Gyroscope - rps");
  display.print("X: ");
  display.print(g.gyro.x, 1);
  display.print(", Y: ");
  display.print(g.gyro.y, 1);
  display.print(", Z: ");
  display.print(g.gyro.z, 1);
  display.println("");
}

void display_max_acc(double max_acc) {
  display.println("Max Acc - m/s^2");
  display.print("Mag: ");
  display.print(max_acc, 1);
  display.println("");
}