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

//Setting up Variables
//OLED setup
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
const byte oled_address = 0x3C; //Got this from clicking question mark on oled

//MPU setup
Adafruit_MPU6050 mpu;
const byte mpu_address = 0x68; //from wokwi queston mark

//Dot Matrix Setup, DIN=51, CLK=52, CS=53, 1 device
LedControl lc = LedControl(51, 52, 53, 1);


//Pin Variables
const int joyYpin = A0;      // Joystick Vertical (Menu Scroll / Depth)
const int joyXpin = A1;      // Joystick Horizontal (Drill Tilt)
const int btnSelect = 2;     // button 1 for select
const int btnBack = 3;       // button 2 for going back, both are connected to interrupt pins
const int stepPin = 4;       // Stepper Motor Step
const int dirPin = 5;        // Stepper Motor Direction
const int servoPin = 6;      // Servo PWM pin
const int redPin = 9;        // pins for the rgb led
const int greenPin = 10;
const int bluePin = 11;
const int buzzerPin = 8;     // for buzzer
const int txLED=7; //For level 3 transmission led

// Doing FSM for the OLED, better than doing if again and again
enum SystemState { MAIN_MENU, CONTROL_DRILL, CRASH_ALERT, VIEW_READING, SEND_DATA};
SystemState currentState = MAIN_MENU;
int menuCursor = 0; // Tracks which menu item is highlighted
int maxItems = 3; //number of menu items

// Interrupt Variables
volatile bool selectTrigger = false;
volatile bool backTrigger = false;

// Motor running variables
Servo drillTilt; //to run servo this object
int currentDepth = 0; // Tracks stepper position

//rock hit variables
bool isJammed = false; //becomes true when threshold cross
float currentJerk = 0.0;
const float jerkThreshold = 4.0;
int jamDepth = 0; //where we jammed the drill

// Timer variables
unsigned long lastScrollTime = 0;
const int scrollDelay = 200; // millisec between menu scrolls

unsigned long lastStepTime = 0;
const int stepInterval = 5; // I put 5 millis gap for a pulse so as to control sensitivity, I can actually see the numbers moving...
//Without a interval for a pulse, the default time to execute a statement( ~ 3 us), is the time for a pulse, this way the numbers increase too fast and u couldn't see them...
bool stepState = LOW; // Tracks the high/low state of the step pin

unsigned long lastDisplayTime = 0;
const int displayInterval = 100; // Refresh OLED every 100 ms

unsigned long crashTimer = 0;

//Struct to store the crash data 
struct HazardData
{
  int depth;
  float jerk;
};
const int MAX_HAZARDS = 8; //Since our dot matrix will display only 8 hazards at a time. I 
HazardData crashLog[MAX_HAZARDS];
int crashCount = 0; //Index variable to our memory array

//Variables for Depth limit on dot matrix
const int STEPS_PER_ROW = 40;        // stepper steps per dot matrix row
const int MAX_DEPTH = 8 * STEPS_PER_ROW; // software limit = 320 steps

// Variable for looking through Readings
int viewCursor = 0;

// Dot matrix animation
unsigned long lastMatrixTime = 0;
const int matrixInterval = 150; //150 ms gap b/w refreshing of dot matrix to see drill moving
int animFrame = 0; // toggles 0/1 for spin effect

//variable to confirm descent
bool isDescending = false;

//variables for transmission data
int maxDepth=0;
int yieldScore=0;
int penaltyFactor = 11; //Multiply the number of rocks hit by 11

//Variables for transmission to LED
bool isSending = false;
int currentBit = 0;         // which of 11 bits we're on (0-10)
byte txPayload = 0;         // the 8-bit score to transmit
byte txParity = 0;          // calculated parity bit
const int timeUnit= 400; //LED will flash for 400 ms in one pulse
unsigned long lastBitTime = 0;
int lastBitSent=0;




// Interrupt for Push Button 1 (Select)
void onSelectPress() {
  selectTrigger = true;
}

// Interrupt for Push Button 2 (Go Back)
void onBackPress() {
  backTrigger = true;
}


void setup() {

  //Pin initialise
  pinMode(joyYpin, INPUT);
  pinMode(joyXpin, INPUT);
  pinMode(btnSelect, INPUT_PULLUP);
  pinMode(btnBack, INPUT_PULLUP);
  pinMode(stepPin, OUTPUT);
  pinMode(dirPin, OUTPUT);
  pinMode(redPin, OUTPUT);
  pinMode(greenPin, OUTPUT);
  pinMode(bluePin, OUTPUT);
  pinMode(buzzerPin, OUTPUT);
  pinMode(txLED, OUTPUT);

  //Attach Servo
  drillTilt.attach(servoPin);

  //Setting up interrupts
  attachInterrupt(digitalPinToInterrupt(btnSelect), onSelectPress, FALLING);
  attachInterrupt(digitalPinToInterrupt(btnBack), onBackPress, FALLING);

  //Start OLED
  //oled will run on 3.3 V to display stuff so I used this function in bracket
  display.begin(SSD1306_SWITCHCAPVCC, oled_address);
  display.clearDisplay();
  display.setTextColor(WHITE);

  //start MPU and set ranges
  mpu.begin(mpu_address);
  mpu.setGyroRange(MPU6050_RANGE_500_DEG);

  //start dot matrix
  lc.shutdown(0, false);   // Wake up the matrix
  lc.setIntensity(0, 8);   // Brightness (0-15) set to medium
  lc.clearDisplay(0);

}

void loop() {
  //checking interrupt
  unsigned long currentMillis = millis();

  if (selectTrigger) {
    selectTrigger = false; // Reset flag
    if (currentState == MAIN_MENU) {

      switch (menuCursor) {
        case 0:
          currentState = CONTROL_DRILL;
          display.clearDisplay();
          break;

        case 1:
          currentState = VIEW_READING;
          viewCursor = 0; //start from first hazard
          display.clearDisplay();
          break;
        
        case 2:
          currentState = SEND_DATA;
          yieldScore= calculateYield();
          startTransmission();
          display.clearDisplay();
          break;
      }
    }
  }

  if (backTrigger) {
    backTrigger = false; // Reset flag
    if (currentState == CONTROL_DRILL || currentState == VIEW_READING || currentState == SEND_DATA) {
      menuCursor=0; // back to top
      currentState = MAIN_MENU;
      display.clearDisplay();
    }
  }

  // Based on current state we call functions
  switch (currentState) {
    case MAIN_MENU:
      analogWrite(redPin, 0); //Turning off the RGB when not drilling
      analogWrite(greenPin, 0);
      analogWrite(bluePin, 0);
      handleMainMenu();
      break;

    case CONTROL_DRILL:
      handleDrillControl();
      drawDrill(isDescending);
      break;

    case CRASH_ALERT:
      // wait for 2 seconds to show the error, then go to menu
      if (currentMillis - crashTimer >= 2000) {
        digitalWrite(buzzerPin, LOW); // Turn off buzzer
        currentState = MAIN_MENU;
      }
      drawDrill(isDescending);
      break;

    case VIEW_READING:
      handleViewReadings();
      drawGraph();
      break;
    
    case SEND_DATA:
      handleSendData();
      runTransmission();
      drawDrill(false); // matrix stays active
      break;
  }
}
//fn. to handle main menu
void handleMainMenu() {
  unsigned long currentMillis = millis();
  int yVal = analogRead(joyYpin);

  //Setting the menucursor to the value needed
  //I assumed 300 and 700 as threshholsds for the joystick to make a little gap from center position, u cn use whatever u wish
  //i.e. if yval>700, it means UP
  if (currentMillis - lastScrollTime >= scrollDelay) {
    if (yVal > 700 && menuCursor > 0) {
      // Joystick pushed UP
      menuCursor--; 
      lastScrollTime = currentMillis;
    } else if (yVal < 300 && menuCursor < maxItems - 1) {
      // Joystick pushed DOWN
      menuCursor++;
      lastScrollTime = currentMillis;
    }
  }

  //Draw the Menu
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println(F(" MAIN MENU ")); //F becus I want it to store in flash memory, saves space on arduino

  // Highlight the selected option
  if (menuCursor == 0)
    display.setTextColor(BLACK, WHITE); // Invert colors to highlight
  else display.setTextColor(WHITE, BLACK);

  display.setCursor(0, 20); //moving the text down for better look
  display.println(F("1. Control Drill"));

  if (menuCursor == 1)
    display.setTextColor(BLACK, WHITE);
  else display.setTextColor(WHITE, BLACK);

  display.setCursor(0, 35); 
  display.println(F("2. View Readings"));
  display.setTextColor(WHITE, BLACK);

  if (menuCursor == 2)
    display.setTextColor(BLACK, WHITE);
  else display.setTextColor(WHITE, BLACK);
  
  display.setCursor(0, 50); 
  display.println(F("3. Send Data"));

  display.setTextColor(WHITE, BLACK); 
  display.display();
}

//fn. to do drill
void handleDrillControl() {
  unsigned long currentMillis = millis();
  int yVal = analogRead(joyYpin);
  int xVal = analogRead(joyXpin);
  String directionStr = "Stationary";
  isDescending = false;
  //Checking for Crash
  isRockHit();
  // Servo Control
  // Map the 0-1023 analog read to a 0-180 degree angle
  int tiltAngle = map(xVal, 0, 1023, 0, 180);
  drillTilt.write(tiltAngle);

  // If we are jammed, check if we are less then jamDepth now and cleared.
  //I coded it as being clear even if we are one step back. u can make a safety margin if u want.
  if (isJammed && currentDepth < jamDepth) {
    isJammed = false; // clear from rock
  }

  // Stepper
  if (yVal > 700) {
    // Retract (Upwards) depth dec.
    digitalWrite(dirPin, LOW);
    directionStr = "Upwards";

    // sending pulse to stepper
    if (currentMillis - lastStepTime >= stepInterval) {
      lastStepTime = currentMillis;
      stepState = !stepState; // Toggle between HIGH and LOW
      digitalWrite(stepPin, stepState);
      if ((stepState == LOW) && (currentDepth > 0)) currentDepth--; // we count depth only when one pulse complete
    }
  }
  else if (yVal < 300) {
    // Descend (Downwards) depth inc.
    isDescending=true;
    if (isJammed) //We can't go down if Jammed
      directionStr = "BLOCKED";

    else if (currentDepth >= MAX_DEPTH) //setting limit in software
      directionStr = "MAX DEPTH";

    else {
      digitalWrite(dirPin, HIGH);
      directionStr = "Downwards";

      if (currentMillis - lastStepTime >= stepInterval) {
        lastStepTime = currentMillis;
        stepState = !stepState;
        digitalWrite(stepPin, stepState);
        if (stepState == LOW){
          currentDepth++; // we count depth only when one pulse complete
          if(currentDepth > maxDepth)  maxDepth= currentDepth; //calculating max depth reached
        }
      }
    }
  }
  else {
    // Joystick is centered
    digitalWrite(stepPin, LOW);
  }

  // Live Display (Update every 100ms so we don't lag the motors)
  if (currentMillis - lastDisplayTime >= displayInterval) {
    lastDisplayTime = currentMillis;
    display.clearDisplay();

    display.setCursor(0, 0);
    display.println(F("--- DRILL CONTROL ---"));

    display.setCursor(0, 20);
    display.print(F("Depth: "));
    display.println(currentDepth);

    display.print(F("Tilt: "));
    display.print(tiltAngle);
    display.println(F(" deg"));

    display.print(F("Dir: "));
    display.println(directionStr);

    // Show warning on oled if currently jammed
    if (isJammed) {
      display.setCursor(0, 50);
      display.println(F("WARNING: DRILL STUCK"));
    }
    display.display();
  }
}

//fn. to check for rock hit
void isRockHit() {
  sensors_event_t Accelerometer, Gyroscope, Temperature;
  mpu.getEvent(&Accelerometer, &Gyroscope, &Temperature);

  float gyro_x = Gyroscope.gyro.x;
  float gyro_y = Gyroscope.gyro.y;

  currentJerk = sqrt((gyro_x * gyro_x) + (gyro_y * gyro_y));

  //mapping the jerk value to 0-255 for LED, multiplying by 100 to get accurate value since map takes int input
  int pwmVal = map(currentJerk * 100, 0, jerkThreshold * 100, 0, 255);

  if (pwmVal >= 255) //if our jerk is over the threshold rn
    pwmVal = 255;

  int redVal = pwmVal;
  int blueVal = 255 - redVal; // values for the LED

  analogWrite(redPin, redVal);
  analogWrite(greenPin, 0);
  analogWrite(bluePin, blueVal);

  if (pwmVal == 255 && !isJammed)
  {

    isJammed = true;
    jamDepth = currentDepth; //remember where we crashed
    currentState = CRASH_ALERT;
    crashTimer = millis();

    if (crashCount < MAX_HAZARDS) { //adding crash to memory
      crashLog[crashCount].depth = jamDepth;
      crashLog[crashCount].jerk = currentJerk;
      crashCount++;
    }

    digitalWrite(buzzerPin, HIGH);

    // Draw the crash screen instantly
    display.clearDisplay();
    display.setTextSize(2); // Make it big
    display.setCursor(0, 10);
    display.println(F("HARD ROCK"));
    display.println(F("  HIT!   "));
    display.display();
    display.setTextSize(1); // Reset for later
  }

}

void handleViewReadings() {
  unsigned long currentMillis = millis();
  int yVal = analogRead(joyYpin);

  // If no hazard logged yet
  if (crashCount == 0) {
    display.clearDisplay();
    display.setCursor(0, 0);
    display.println(F("  VIEW READINGS  "));
    display.setCursor(0, 25);
    display.println(F("No hazards logged."));
    display.display();
    return; // nothing else to do
  }

  // Scroll through hazards
  if (currentMillis - lastScrollTime >= scrollDelay) {
    if (yVal > 700 && viewCursor > 0) {
      viewCursor--;
      lastScrollTime = currentMillis;
    }
    else if (yVal < 300 && viewCursor < crashCount - 1) {
      viewCursor++;
      lastScrollTime = currentMillis;
    }
  }

  // Display current hazard
  if (currentMillis - lastDisplayTime >= displayInterval) {
    lastDisplayTime = currentMillis;
    display.clearDisplay();

    display.setCursor(0, 0);
    display.println(F("  VIEW READINGS  "));

    display.setCursor(0, 15);
    display.print(F("Hazard: "));
    display.print(viewCursor + 1); // 1-indexed for readability
    display.print(F(" / "));
    display.println(crashCount);

    display.setCursor(0, 30);
    display.print(F("Depth: "));
    display.println(crashLog[viewCursor].depth);

    display.setCursor(0, 45);
    display.print(F("Jerk:  "));
    display.println(crashLog[viewCursor].jerk);

    display.display();
  }
}

byte pattern[] = { //pattern to print on LED as per distance from tip, 0 index means tip itself
  0x18,  // 0 0 0 1 1 0 0 0  distance 0
  0x18,  // 0 0 0 1 1 0 0 0  distance 1
  0x3C,  // 0 0 1 1 1 1 0 0  distance 2
  0x3C,  // 0 0 1 1 1 1 0 0  distance 3
  0x7E,  // 0 1 1 1 1 1 1 0  distance 4
  0x7E,  // same for fartherr rows
  0x7E,
  0x7E,  // 0 1 1 1 1 1 1 0
};

// Tip animation rows of bits, to show moving
byte tipFrame0 = 0x10; // 0 0 0 1 0 0 0 0
byte tipFrame1 = 0x08; // 0 0 0 0 1 0 0 0

void drawDrill(bool isDescending) {

  unsigned long currentMillis = millis();
  if (currentMillis - lastMatrixTime < matrixInterval) return;
  lastMatrixTime = currentMillis;

  // Only animate tip if motor is actively descending
  if (isDescending)
    animFrame = !animFrame; // toggle only while moving

  int drillRow = currentDepth / STEPS_PER_ROW;
  lc.clearDisplay(0);

  if(drillRow>=7) drillRow=7; //From 0-7 only
  
  // Draw the permanent cone trail before drill tip
  for (int r = 0; r < drillRow && r <= 7; r++)
    lc.setRow(0, r, pattern[drillRow-r]);

  // Draw the animated tip at drillRow
  if (drillRow <= 7) {
    byte tip;
    if(isDescending)
    tip = (animFrame) ? tipFrame1 : tipFrame0;
    else
    tip =0x18;
    lc.setRow(0, drillRow, tip);
  }

  // Draw hazard dots on top of everything (column 7, leftmost)
  for (int i = 0; i < crashCount; i++) {
    int hazardRow = crashLog[i].depth / STEPS_PER_ROW;
    lc.setLed(0, hazardRow, 7, true);
  }
}

void drawGraph() {
  unsigned long currentMillis = millis();
  if(crashCount==0) //Nothing to draw
  {
    lc.clearDisplay(0);
    return;
  }

  if (currentMillis - lastMatrixTime < matrixInterval) return; //Add check to draw graph once every 150 ms
  lastMatrixTime = currentMillis;

  lc.clearDisplay(0);
  int i;
  //Finding max and min jerk to plot based on range.
  float minJerk = crashLog[0].jerk;
  float maxJerk = crashLog[0].jerk;

  for(i = 0; i < crashCount; i++)
  {
    if (crashLog[i].jerk > maxJerk)
      maxJerk = crashLog[i].jerk;

    if (crashLog[i].jerk < minJerk)
      minJerk = crashLog[i].jerk;
  }

  for(i = 0; i < crashCount; i++)
  { int barHeight;
    if (maxJerk == minJerk)
      barHeight = 8; // all same jerk, show full bars

    //Again multiplying by 100 to adjust for integer input, and also mapping based on range now.
    //from 1-8 because atleast 1 row per hazard.
    else
      barHeight = map((crashLog[i].jerk - minJerk) * 100, 0, (maxJerk - minJerk) * 100, 1, 8);

    //Lighting up the column acc. to hazard
    for(int r = 7; r >= 8 - barHeight; r--)
      lc.setLed(0, r, 7-i, true);//rightmost column is zero thats why

  }
  // Highlight currently viewed hazard
    lc.setLed(0, 0, 7-viewCursor, true);//rightmost column is zero
}

int calculateYield() {
  int score = map(maxDepth, 0 , MAX_DEPTH, 0, 100);
  score -= penaltyFactor*crashCount;
  return (score<0) ? 0 : score;
}

void handleSendData() {

  unsigned long currentMillis= millis();

  if (currentMillis - lastDisplayTime >= displayInterval) {
    lastDisplayTime = currentMillis;
    display.clearDisplay();

    display.setCursor(0, 0);
    display.println(F("  SEND DATA  "));

    display.setCursor(0, 20);
    display.print(F("Yield Score: "));
    display.println(yieldScore);

    display.setCursor(0, 35);
    if(isSending)
    display.println(F("Transmitting Data..."));
    else display.println(F("Transmission Done"));


    display.setCursor(0, 50);
    display.print(F("Bit: "));
    display.print((isSending) ? (currentBit+1) : (lastBitSent+1));
    display.println(F(" / 11"));

    display.display();
  }

  


}

//To calculate payload data before starting
void startTransmission() {
  //calculating the payload data here only
  txPayload = (byte)yieldScore;
  // calculate parity
  txParity = 0;
  for (int i = 0; i < 8; i++)
    txParity ^= (txPayload >> i) & 1; //calculating parity by taking the unit bit and XORing it again and again
  currentBit = 0;
  isSending = true;
  lastBitTime = millis();
}

void runTransmission() {
  if (!isSending) return;

  unsigned long currentMillis = millis();
  if (currentMillis - lastBitTime < timeUnit) return;
  lastBitTime = currentMillis;

  if (currentBit == 0)
    digitalWrite(txLED, HIGH); // start bit
  else if (currentBit >= 1 && currentBit <= 8) {
    int bitIndex = 8 - currentBit; //Need to send MSB first
    bool bitVal = (txPayload >> bitIndex) & 1; //Taking out bits from MSB to LSB
    digitalWrite(txLED, bitVal);
    }
  else if (currentBit == 9)
    digitalWrite(txLED, txParity);

  else if (currentBit == 10) {
    digitalWrite(txLED, LOW); // stop bit
    isSending = false;
    lastBitSent= 10;
    currentBit = 0;
    return;
    }
  currentBit++;
}