#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Servo.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>
#include <math.h>
#define STEP_PIN 6
#define DIR_PIN 7
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define RED_PIN   10
#define GREEN_PIN 11
#define BLUE_PIN  12
#define BUZZER_PIN 8

Adafruit_SSD1306 display(
  SCREEN_WIDTH,
  SCREEN_HEIGHT,
  &Wire,
  -1
);

Servo drillServo;
Adafruit_MPU6050 mpu;

enum State
{
  MAIN_MENU,
  CONTROL_DRILL,
  HARD_ROCK_SCREEN,
  DRILL_STUCK_SCREEN
};

State currentState = MAIN_MENU;

int previousX = -1;
int previousTilt = -1;
float previousJerk = -1;

volatile bool selectPressed = false;
volatile bool backPressed = false;

int depth = 0;
int tiltAngle = 90;
String direction = "Stationary";

sensors_event_t a, g, temp;
float gyroX;
float gyroY;
float jerk = 0;
const float HARD_THRESHOLD = 2.0;
const float MAX_THRESHOLD  = 5.0;

bool drillJammed = false;
int hazardStartDepth = -1;
bool insideHazardZone = false;

enum RockState
{
    SAFE_ZONE,
    HAZARD_ZONE,
    JAM_ZONE
};
RockState currentRockState = SAFE_ZONE;


void stepMotor()
{
  digitalWrite(STEP_PIN, HIGH);
  delayMicroseconds(500);
  digitalWrite(STEP_PIN, LOW);
  delayMicroseconds(500);
}

void selectISR()
{
  selectPressed = true;
}

void backISR()
{
  backPressed = true;
}

void setup()
{
  if(!display.begin(
        SSD1306_SWITCHCAPVCC,
        0x3C))
  {
    while(true);
  }

  pinMode(STEP_PIN, OUTPUT);
  pinMode(DIR_PIN, OUTPUT);
  digitalWrite(DIR_PIN, HIGH);
  Serial.begin(115200);

  if (!mpu.begin())
  {
    Serial.println("MPU6050 not found!");

    while (true)
    {
    }
  }

  pinMode(2, INPUT_PULLUP);
  pinMode(3, INPUT_PULLUP);

  attachInterrupt(
    digitalPinToInterrupt(2),
    selectISR,
    FALLING
  );

  attachInterrupt(
    digitalPinToInterrupt(3),
    backISR,
    FALLING
  );

  drillServo.attach(9);

  pinMode(RED_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(BLUE_PIN, OUTPUT);

  pinMode(BUZZER_PIN, OUTPUT);
}

void loop()
{
  int x = analogRead(A0);
  int y = analogRead(A1);

  if (x != previousX || tiltAngle != previousTilt)
  {
    Serial.print("X = ");
    Serial.print(x);

    Serial.print("  Tilt = ");
    Serial.println(tiltAngle);

    previousX = x;
    previousTilt = tiltAngle;
  }

  if(x < 300)
  {
    tiltAngle--;
  }
  if(x > 700)
  {
    tiltAngle++;
  }
  if(y < 300 && !drillJammed)
  {
    digitalWrite(DIR_PIN, HIGH);
    stepMotor();
    depth++;
    direction = "Down";
  }

  if(y > 700)
  {
    digitalWrite(DIR_PIN, LOW);
    stepMotor();
    if(depth > 0)
    {
      depth--;
    }
    if (drillJammed && depth <= hazardStartDepth && jerk < HARD_THRESHOLD)
    {
      drillJammed = false;
      noTone(BUZZER_PIN);
      insideHazardZone = false;
      Serial.println("JAM CLEARED");
      currentState = CONTROL_DRILL;
    }
    direction = "Up";
  }

  if(y >= 300 && y <= 700)
  {
    direction = "Stationary";
  }

  tiltAngle = constrain(
  tiltAngle,
  0,
  180
  );

  drillServo.write(tiltAngle);

  mpu.getEvent(&a, &g, &temp);
  gyroX = g.gyro.x;
  gyroY = g.gyro.y;
  float jerk = sqrt(
    gyroX * gyroX +
    gyroY * gyroY
  );

 RockState newRockState;

if (jerk > MAX_THRESHOLD)
{
    newRockState = JAM_ZONE;
}
else if (jerk > HARD_THRESHOLD)
{
    newRockState = HAZARD_ZONE;
}
else
{
    newRockState = SAFE_ZONE;
}

if (newRockState != currentRockState)
{
    currentRockState = newRockState;

    switch (currentRockState)
    {
        case SAFE_ZONE:
            Serial.println("SAFE");
            break;

        case HAZARD_ZONE:
            Serial.println("HARD ROCK");
            break;

        case JAM_ZONE:
            Serial.println("MAX RED");
            break;
    }
}

  if (jerk != previousJerk)
  {
      Serial.print("Jerk = ");
      Serial.println(jerk);

      previousJerk = jerk;
  }

  if(jerk > HARD_THRESHOLD && !insideHazardZone)
  {
      insideHazardZone = true;
      hazardStartDepth = depth;
      Serial.println("Entered Hazard Zone");
      Serial.print("Hazard Depth Stored = ");
      Serial.println(depth);
  }

  if(jerk <= HARD_THRESHOLD)
  {
      insideHazardZone = false;
  }

  if (jerk > MAX_THRESHOLD && !drillJammed)
  {
      drillJammed = true;
      tone(BUZZER_PIN, 1000);
      currentState = HARD_ROCK_SCREEN;
      Serial.println("HARD ROCK HIT!");
      Serial.print("Hazard Start Depth = ");
      Serial.println(hazardStartDepth);
      Serial.print("Current Depth = ");
      Serial.println(depth);
      Serial.print("Jam Depth = ");
      Serial.println(depth);
  }

  if(selectPressed)
  {
    selectPressed = false;

     if(currentState == MAIN_MENU)
      {
        if(drillJammed)
        {
          currentState = DRILL_STUCK_SCREEN;
        }
        else
        {
          currentState = CONTROL_DRILL;
        }
      }
  }


  if(backPressed)
  {
    backPressed = false;
    currentState = MAIN_MENU;
  }


  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  if(currentState == MAIN_MENU)
  {
    display.setCursor(20,0);
    display.println("MAIN MENU");
    display.setCursor(10,25);
    display.println("> Control Drill");
  }
  if(currentState == CONTROL_DRILL)
  {
    display.setCursor(0,0);
    display.println("CONTROL DRILL");
    display.setCursor(0,20);
    display.print("Depth: ");
    display.println(depth);
    display.setCursor(0,35);
    display.print("Tilt: ");
    display.println(tiltAngle);
    display.setCursor(0,50);
    display.print("Dir: ");
    display.println(direction);
  }
  if(currentState == HARD_ROCK_SCREEN)
  {
    display.clearDisplay();
    display.setTextSize(2);
    display.setCursor(0,20);
    display.println("HARD");
    display.println("ROCK!");
    display.display();
  }

  if(currentState == DRILL_STUCK_SCREEN)
  {
    display.clearDisplay();

    display.setTextSize(1);

    display.setCursor(0,15);
    display.println("WARNING:");

    display.setCursor(0,30);
    display.println("DRILL STUCK");

    display.display();
  }

  display.display();

  if (jerk < 2)
  {
      analogWrite(RED_PIN, 0);
      analogWrite(GREEN_PIN, 0);
      analogWrite(BLUE_PIN, 255);
  }
  else if (jerk < 5)
  {
      analogWrite(RED_PIN, 255);
      analogWrite(GREEN_PIN, 0);
      analogWrite(BLUE_PIN, 255);
  }
  else
  {
      analogWrite(RED_PIN, 255);
      analogWrite(GREEN_PIN, 0);
      analogWrite(BLUE_PIN, 0);
  }
}