#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>

// Pin definitions for I2C
#define SDA_PIN 18
#define SCL_PIN 19

// Screen dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Create the OLED display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Initialize the MPU6050
Adafruit_MPU6050 mpu;

void setup() {
  Serial.begin(115200);
  Wire.begin(SDA_PIN, SCL_PIN);

  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x64
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  display.display();
  delay(2000);
  display.clearDisplay();

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

  mpu.setAccelerometerRange(MPU6050_RANGE_8_G);
  mpu.setGyroRange(MPU6050_RANGE_500_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);
}

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

  float roll = atan2(a.acceleration.y, a.acceleration.z) * 180 / PI;
  float pitch = atan2(-a.acceleration.x, sqrt(a.acceleration.y * a.acceleration.y + a.acceleration.z * a.acceleration.z)) * 180 / PI;

  display.clearDisplay();
  drawArtificialHorizon(roll, pitch);
  display.display();

  delay(100);
}

void drawArtificialHorizon(float roll, float pitch) {
  const int lineLength = SCREEN_WIDTH; // max line length across the display
  const int horizonY = SCREEN_HEIGHT / 2; // center of the screen
  const float radians = -roll * PI / 180;

  // Calculate the endpoints of the horizon line
  float x1 = horizonY * tan(radians);
  float y1 = horizonY;
  float x2 = x1 + lineLength * cos(radians);
  float y2 = y1 - lineLength * sin(radians);

  // Set the color to white
  display.setTextColor(SSD1306_WHITE);

  // Draw the sky (upper part)
  display.fillRect(0, 0, SCREEN_WIDTH, horizonY, SSD1306_WHITE);

  // Draw the ground (lower part)
  display.fillRect(0, horizonY, SCREEN_WIDTH, SCREEN_HEIGHT - horizonY, SSD1306_BLACK);

  // Draw the horizon line
  display.drawLine(x1, y1, x2, y2, SSD1306_BLACK);

  // Draw roll indicator
  display.setCursor(SCREEN_WIDTH / 2 - 12, horizonY - 7);
  display.print((int)roll);
}