// header files - tried to keep it to a minimum
// got rid of stepper, servo, mpu sensor libraries
//
#include <Wire.h> // I2C for OLED and MPU6050 
#include <Adafruit_GFX.h> // graphics for OLED
#include <Adafruit_SSD1306.h> // OLED
#include <LedControl.h> // 8x8 dot matrix

// PINS, CONST_VALUES and variableNames

// pin definitions for reference
const int JOY_X = A1;         
const int JOY_Y = A0;         
const int BTN_SELECT = 3;     
const int BTN_BACK = 2;       
const int TX_LED_PIN = 4;     
const int SERVO_PIN = 5;      
const int BUZZER_PIN = 12;    
const int RGB_BLUE = 44;      
const int RGB_GREEN = 46;     
const int RGB_RED = 45;       
const int STEPPER_PINS[4] = {10, 11, 9, 8}; // in the order A+ A- B+ B-

// OLED display object
Adafruit_SSD1306 display(128, 64, &Wire, -1);

// Dot matrix display object
LedControl lc = LedControl(51, 52, 53, 1);

// stepper motor var.s
// 
int stepIndex = 0; 
// NEW
const bool stepSequence[4][4] = {
  {HIGH, LOW,  HIGH, LOW}, // top right position
  {LOW,  HIGH, HIGH, LOW}, // bottom right position
  {LOW,  HIGH, LOW,  HIGH}, // bottom left position
  {HIGH, LOW,  LOW,  HIGH} // top left position
};

// MPU6050 var.s 
const int MPU_ADDR = 0x68; //mpu default address 
int16_t gyro_x, gyro_y; //global raw jerk variable
float currentJerk = 0.0;
const float MAX_JERK = 100.0; //jerk threshold

// drill var.s 
int depth = 0;
int maxDepth = 0; 
int tilt_angle = 90; 
String dir = "Stationary";

//FSM - statess
//
enum MachineState {
  STATE_MENU, // state 0 - displays menu
  STATE_DRILL, // state 1 - controls drill, displays telemetry
  STATE_CRASH_SCREEN, // for 2 seconds right after high jerk
  STATE_JAMMED, // in control drill state when jammed
  STATE_VIEW_LOGS, // state 2 - view logs and graph
  STATE_TRANSMIT // state 3 -  excavation yield score transmission    
};
volatile MachineState systemState = STATE_MENU; //volatile since we modify inside select and back ISR
MachineState lastState = STATE_MENU; 

//state switching related varibales
volatile int menuCursor = 1;  
volatile bool selectPressed = false;
volatile bool backPressed = false;
volatile bool isJammed = false;
volatile unsigned long crashTime = 0;
int retractCount = 0; 

//safety var.s
struct Hazard { //rock hard hit log data
  int depth;
  float jerk;
  int matrix_x; //drill tip
  int matrix_y;
};
Hazard hazardLogs[8]; 
int hazardCount = 0; // out of the max 8
int logIndex = 0;    

// transmisison var.s
int yieldScore = 0;
bool txData[11]; // 11 total bits in one data packet
int txBitIndex = 0; // to traverse through
unsigned long txPreviousMillis = 0; // control timings
const long txInterval = 200; //1 bit every 200 ms

// non blocking timing related var.s
unsigned long previousMillis = 0;
const long interval = 50; 
volatile unsigned long lastSelectTime = 0;
volatile unsigned long lastBackTime = 0;
const unsigned long debounceDelay = 100; 


// USER DEFINED FUNCTIONS

// setting angle for servo without lib.
//
void setServo(int angle) {
  if (angle < 0) angle = 0;
  if (angle > 180) angle = 180;
  OCR3A = map(angle, 0, 180, 1088, 4800); 
  // OCR3A controls pin 5 PWM's by setting the trigger (when high to low transition happens in the total 20ms it takes to count from 0 to 39999)
}

// moving stepper one step without lib.
//
void setStepMotor(int dir_step) {
  stepIndex += dir_step;
  // stepindex to controol if we want to rot forward or reverse
  if (stepIndex > 3) stepIndex = 0;
  if (stepIndex < 0) stepIndex = 3;
  for (int i = 0; i < 4; i++) {
    digitalWrite(STEPPER_PINS[i], stepSequence[stepIndex][i]);
  }
}

// reading MPU6050 directly without lib.
void readGyroMPU() {
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x43); // the address with the gyro data as per datasheet
  Wire.endTransmission(false); // keep the connection open
  Wire.requestFrom(MPU_ADDR, 4, true); // requesting 4 bytes of data from mpu, starting from the address 0x43 
  gyro_x = (Wire.read() << 8 | Wire.read()); // read 2 bytes, shift the first one by 8 bytes and combine
  gyro_y = (Wire.read() << 8 | Wire.read()); // same for Y, the enxt subsequent address is automatically read since we requested 4 bytes
}

// check if valid position and light up in Dot Matrix
void safeDraw(int x, int y, bool state) {
  if (x >= 0 && x <= 7 && y >= 0 && y <= 7) {
    lc.setLed(0, y, x, state);
  }
}

// draw the drill shape, according to depth & angle
void drawDrill(int base_x, int base_y, int angle, bool isWide) {
  lc.clearDisplay(0);
  
  // display previous high jerk pts
  for(int i = 0; i < hazardCount; i++) safeDraw(hazardLogs[i].matrix_x, hazardLogs[i].matrix_y, true); 

  //isWide controls its current state (out of 2) to simulate like an animation 
  if (angle <= 60) { // LEFT 
    if (isWide) {
      // column of 5
      safeDraw(base_x, base_y-2, true); 
      safeDraw(base_x, base_y-1, true); 
      safeDraw(base_x, base_y, true); 
      safeDraw(base_x, base_y+1, true); 
      safeDraw(base_x, base_y+2, true);
      // column of 3 to its left
      safeDraw(base_x-1, base_y-1, true); 
      safeDraw(base_x-1, base_y, true); 
      safeDraw(base_x-1, base_y+1, true); 
      //coumn of 1 to its left left
      safeDraw(base_x-2, base_y, true); 
    } else {
      //column of 3
      safeDraw(base_x, base_y-1, true); 
      safeDraw(base_x, base_y, true); 
      safeDraw(base_x, base_y+1, true);
      // collum of 1 to its left and then another to its left 
      safeDraw(base_x-1, base_y, true); 
      safeDraw(base_x-2, base_y, true); 
    }
  } 
  else if (angle >= 120) { // RIGHT 
    if (isWide) {
      // column of 5
      safeDraw(base_x, base_y-2, true); 
      safeDraw(base_x, base_y-1, true);
      safeDraw(base_x, base_y, true); 
      safeDraw(base_x, base_y+1, true);
      safeDraw(base_x, base_y+2, true);
      // column of 3 to its right
      safeDraw(base_x+1, base_y-1, true);
      safeDraw(base_x+1, base_y, true); 
      safeDraw(base_x+1, base_y+1, true); 
      //coumn of 1 to its right right
      safeDraw(base_x+2, base_y, true); 
    } else {
      //column of 3 
      safeDraw(base_x, base_y-1, true);
      safeDraw(base_x, base_y, true); 
      safeDraw(base_x, base_y+1, true);
      // collum of 1 to its right and then another to its right 
      safeDraw(base_x+1, base_y, true);
      safeDraw(base_x+2, base_y, true); 
    }
  } 
  else { // DOWN 
    if (isWide) { // all horizontal wise
      //first row 5 
      safeDraw(base_x-2, base_y, true);
      safeDraw(base_x-1, base_y, true); 
      safeDraw(base_x, base_y, true);
      safeDraw(base_x+1, base_y, true); 
      safeDraw(base_x+2, base_y, true);
      // second row 3
      safeDraw(base_x-1, base_y+1, true);
      safeDraw(base_x, base_y+1, true);
      safeDraw(base_x+1, base_y+1, true); 
      // third row 1
      safeDraw(base_x, base_y+2, true); 
    } else {
      //first row 3
      safeDraw(base_x-1, base_y, true);
      safeDraw(base_x, base_y, true);
      safeDraw(base_x+1, base_y, true);
      // second and third row 1
      safeDraw(base_x, base_y+1, true);
      safeDraw(base_x, base_y+2, true); 
    }
  }
}

// do as minimum work as possible in ISR, so only changing flags (fast process) 
// select option ISR
ISR(INT5_vect) {
  unsigned long currentTime = millis();
  if ((currentTime - lastSelectTime) > debounceDelay) {
    selectPressed = true; 
    lastSelectTime = currentTime;
  }
}

// back to main menu ISR
ISR(INT4_vect) {
  unsigned long currentTime = millis();
  if ((currentTime - lastBackTime) > debounceDelay) {
    backPressed = true; 
    lastBackTime = currentTime;
  }
}

//SETUP

void setup() {
  // I2C comm.s
  Wire.begin();

  //pinModes

  //input_pullups for buttons
  //port E bit 4,5 are pins 2,3 (pinout sheet of ard. mega)
  DDRE &= ~((1 << 5) | (1 << 4)); //DDR 0 - input mode
  PORTE |= (1 << 5) | (1 << 4); //PORT 1 when DDR 0 - activates int pullup resistor

  //outputs : we have DDR as 1
  DDRG |= (1 << 5); // Port G, bit 5 - pin 4. (forr TX_LED)
  DDRE |= (1 << 3); // PE3 - pin 5  (for servo)
  DDRB |= (1 << 6); // PB6 - pin 12 (buzzer)
  DDRL |= (1 << 3) | (1 << 4) | (1 << 5); //PL3, PL4, PL5 are pins 46,45,44 (RGB)
  DDRB |= (1 << 5) | (1 << 4); // PB5 = pin 11, PB4 = pin 10 (stepper)
  DDRH |= (1 << 6) | (1 << 5); // PH6 = 9, PH5 = 8 (stepper)
  
  // RGB LED Timers
  TCCR5A = (1 << COM5A1) | (1 << COM5B1) | (1 << COM5C1) | (1 << WGM50); 
  TCCR5B = (1 << CS51) | (1 << CS50);
  
  // for servo
  // TCCR - timer/counter control register controls PWM on pins 
  // since only COM3A1 high (COM3A0 low), its non inverting mode, timer starts HIGH, and after trigger LOW
  // ssetting of trigger will be done in execution dependign on the angle to be set (ratio of high to low)
  TCCR3A = (1 << COM3A1) | (1 << WGM31);
  // CS31 sets a prescaler of 8
  // WGM31, 32, 33 high and WGM30 low - we selecting mode 14 of many, fast pwm, couting upto ICR3 and reset
  TCCR3B = (1 << WGM33) | (1 << WGM32) | (1 << CS31); 
  ICR3 = 39999; // 16Mhz (original freq) / prescaler
  setServo(tilt_angle); // initial 

  // initialise MPU
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x6B);  // power management register of mpu
  Wire.write(0x00);  // write 0 to wake up
  Wire.endTransmission(true);

  // initialise OLED and Dot Matrix
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) for(;;); //endless loop if error
  lc.shutdown(0, false); // 0 - do all operations to the first OLED you find. shutdown - false: wake up
  lc.setIntensity(0, 8);
  lc.clearDisplay(0);

  // attach interrupts
  EICRB |= (1 << ISC41) | (1 << ISC51); 
  //External interrupt Control Register controls how the trigger is triggered, herre during FALLING edge 
  EIMSK |= (1 << INT4) | (1 << INT5); 
  // External interrupt mask register, to say that we're using INT4, INT5 = pin 2,3 as interrupts
}

// LOOP

void loop() {
  // assign as "currentTime" for all the upcoming porcesses
  unsigned long currentMillis = millis();

  // transmission task if in state transit 
  // outside 50 ms interval because timing issues
  if (systemState == STATE_TRANSMIT) {
    if (currentMillis - txPreviousMillis >= txInterval) {
      txPreviousMillis = currentMillis;
      txBitIndex++;
      if (txBitIndex > 10) txBitIndex = 0; 
    }
    //digitalWrite(TX_LED_PIN, txData[txBitIndex]); 

    //longer code but shorter execution because no digitalWrite
    if (txData[txBitIndex]) {
      PORTG |= (1 << 5);  // if txData[txBitIndex] high, set PG5 high
    } else {
      PORTG &= ~(1 << 5); // else, set PG5 low
    }
  } else {
    PORTG &= ~(1 << 5);   // set PG5 (tx led, pin 4) low 
  }
  // loop still technically many many thousands of times per second, accurately, 
  // but our stuff exectes once every 50 ms only, rate limiter without delay
  if (currentMillis - previousMillis >= interval) {
    previousMillis = currentMillis;

    int joy_x_val = analogRead(JOY_X);
    int joy_y_val = analogRead(JOY_Y);

    int red_ill = 0;
    int blue_ill = 0;

    // selecting the option 
    if (selectPressed) {
      selectPressed = false; 
      if (systemState == STATE_MENU) { // select the option if in menu
        if (menuCursor == 1) systemState = isJammed ? STATE_JAMMED : STATE_DRILL;
        else if (menuCursor == 2) systemState = STATE_VIEW_LOGS;
        else if (menuCursor == 3) {
          
          yieldScore = (int)(0.5 * abs(maxDepth)) - (10 * hazardCount);
          if (yieldScore < 0) yieldScore = 0;
          if (yieldScore > 100) yieldScore = 100;

          //here only storing happens, no transmission in led
          txData[0] = LOW; // start bit, 1 bit
          int onesCount = 0;
          for (int i = 0; i < 8; i++) {
            txData[i + 1] = bitRead(yieldScore, i); // actual data, read in binary to 8 bits, i+1 cuz of the start bit extra
            if (txData[i + 1] == HIGH) onesCount++;
          }
          txData[9] = (onesCount % 2 == 0) ? LOW : HIGH; // parity, 1 bit
          txData[10] = HIGH; // stop or end bit, 1 bit 
          txBitIndex = 0; //reset
          systemState = STATE_TRANSMIT; // assign everything then go to tranmistting state
        }
      }
    }
    // returning to main menu
    if (backPressed) {
      backPressed = false; 
      if (systemState == STATE_DRILL || systemState == STATE_JAMMED || systemState == STATE_VIEW_LOGS || systemState == STATE_TRANSMIT) {
        systemState = STATE_MENU;
      }
    }

    // here comes the FINITE STATE MACHINE
    switch (systemState) {

      // sttae 0 - main menu
      case STATE_MENU: 
        
        // set selection cursor up and down, later "move"
        if (joy_y_val < (400)) { menuCursor++; if(menuCursor > 3) menuCursor = 3; }
        else if (joy_y_val > (600)) { menuCursor--; if(menuCursor < 1) menuCursor = 1; }
        
        //display menu
        display.clearDisplay();
        display.setTextSize(1.5);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("Main Menu: ");
        display.setCursor(0, 16);
        if (menuCursor == 1) display.print("> "); else display.print("  ");
        display.println("1. Control Drill");
        display.setCursor(0, 32);
        if (menuCursor == 2) display.print("> "); else display.print("  ");
        display.println("2. View Readings");
        display.setCursor(0, 48);
        if (menuCursor == 3) display.print("> "); else display.print("  ");
        display.println("3. Send Data");
        display.display();

        break;

      // state 1 - control drill
      case STATE_DRILL: {
        // calc jerk
        readGyroMPU(); //reads gyro_x and gyro_y
        currentJerk = sqrt(pow(gyro_x / 131.0, 2) + pow(gyro_y / 131.0, 2));

        //choose red adn blue pwm values and light up based on hot or cold (jerk)
        red_ill = constrain(map((int)currentJerk, 0, MAX_JERK, 0, 255), 0, 255);
        blue_ill = constrain(map((int)currentJerk, 0, MAX_JERK, 255, 0), 0, 255);
        OCR5B = red_ill;  // controls pin 45's PWM
        OCR5C = blue_ill; // pin 44s PWM
        OCR5A = 0;        // pin 46's PWM

        // for displaying drill
        int base_y = map(depth, 200, -200, 0, 7);
        int tip_x = 4;
        int tip_y = base_y; 
        //tip x and tip y are for logging jerk pts
        //
        
        if (tilt_angle <= 60) tip_x = 2;       
        else if (tilt_angle >= 120) tip_x = 6;   
        else tip_y = base_y + 2;                    

        // for logging in case of rock hard hit, set flag, buzzer, change state, turn off stepper coils
        if (currentJerk >= MAX_JERK) {
          if (hazardCount < 8) {
            hazardLogs[hazardCount].depth = depth;
            hazardLogs[hazardCount].jerk = currentJerk;
            hazardLogs[hazardCount].matrix_x = tip_x;
            hazardLogs[hazardCount].matrix_y = tip_y;
            hazardCount++;
          }
          isJammed = true;
          crashTime = millis();
          systemState = STATE_CRASH_SCREEN;
          tone(BUZZER_PIN, 1000); 
          for (int i = 0; i < 4; i++) digitalWrite(STEPPER_PINS[i], LOW); // coil not needed to be on 
          break; 
        }

        //move stepper
        if (joy_y_val < (400)) {
          if (depth > -200) {
            depth--;
            if (depth < maxDepth) maxDepth = depth; 
            dir = "Downwards";
            setStepMotor(-1);
          } else dir = "Max Depth";
        }
        else if (joy_y_val > (600)) {
          if (depth < 200) {
            depth++;
            dir = "Upwards";
            setStepMotor(1);
          } else dir = "Max Height";
        }
        else {
          dir = "Stationary";
          for (int i = 0; i < 4; i++) digitalWrite(STEPPER_PINS[i], LOW); 
        }

        //move servo
        if ((joy_x_val > (600)) && (tilt_angle > 0)) tilt_angle--; 
        else if ((joy_x_val < (400)) && (tilt_angle < 180)) tilt_angle++; 
        setServo(tilt_angle);
        
        // display drill (animated)
        bool spinningFrame = (currentMillis / 150) % 2 == 0; 
        drawDrill(4, base_y, tilt_angle, spinningFrame); 

        // display telemetry dashboard
        display.clearDisplay();
        display.setTextSize(1.5); 
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("Telemetry Display:\n");
        display.println("Depth: " + String(depth));
        display.println("Tilt: " + String(tilt_angle) + " deg");
        display.println("Status: " + dir);
        display.println("Jerk: " + String(currentJerk));
        display.display();
        break;
      }

      // crash screen - 2 seconds then back to menu
      case STATE_CRASH_SCREEN:
        display.clearDisplay();
        display.setTextSize(1.5);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 20);
        display.println("HARD ROCK HIT!");
        display.display();
        if (currentMillis - crashTime > 2000) systemState = STATE_MENU;
        break;

      // inside controll drill ie state 1 when still jammed, hard rock hit, poor driving skills :(
      case STATE_JAMMED: {
        // moving up 10 steps to un-jam
        if (joy_y_val > (600) && depth < 200) { 
          depth++;
          dir = "Retracting";
          setStepMotor(1);
          retractCount++;
          if (retractCount > 10) {
            isJammed = false;
            retractCount = 0;
            noTone(BUZZER_PIN);          
            OCR5B = 0; // Turn off Red PWM 
            systemState = STATE_DRILL;   
          }
        } else { // if not moving up, then jam it
          dir = "JAMMED";
          for (int i = 0; i < 4; i++) digitalWrite(STEPPER_PINS[i], LOW);
        }

        // display drill in dot matrix
        int jam_y = map(depth, 200, -200, 0, 7);
        drawDrill(4, jam_y, tilt_angle, true);

        // warnin and retract in oled
        display.clearDisplay();
        display.setTextSize(1);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("Warning: Drill stuck");
        display.println("");
        display.println("Pull up to retract");
        display.println("Depth: " + String(depth));
        display.display();
        break;
      }

      //state 2 - view logs   
      case STATE_VIEW_LOGS: {
        //move pages up and down, with limit
        if (joy_y_val < (400) && hazardCount > 0) { 
          logIndex++;
          if (logIndex >= hazardCount) logIndex = hazardCount - 1;
        } 
        else if (joy_y_val > (600) && hazardCount > 0) { 
          logIndex--;
          if (logIndex < 0) logIndex = 0;
        }

        // display stored logs in oled
        display.clearDisplay();
        display.setTextSize(1);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("Hazards till now:");
        if (hazardCount == 0) display.println("\nNo hazards recorded.");
        else {
          display.println("\nLog: " + String(logIndex + 1) + " / " + String(hazardCount));
          display.println("Depth: " + String(hazardLogs[logIndex].depth));
          display.println("Jerk: " + String(hazardLogs[logIndex].jerk));
        }
        display.display();

        //display bar graph in dot matrix
        lc.clearDisplay(0);
        for (int i = 0; i < hazardCount; i++) {
          int barHeight = map((int)hazardLogs[i].jerk, MAX_JERK, MAX_JERK + 254, 1, 8); // maximum jerk 354
          barHeight = constrain(barHeight, 1, 8);
          
          //draw graph
          int column = 7 - i;
          for (int r = 7; r > 7 - barHeight; r--) lc.setLed(0, r, column, true); // for for each column
          if (i == logIndex && (currentMillis / 250) % 2 == 0) { //blink the currently viewing log for identification
            for (int r = 0; r < 8; r++) lc.setLed(0, r, column, false);
          }
        }
        break;
      }
      
      // state 3 - transmission of score
      case STATE_TRANSMIT: {
        // OLED Graphic
        display.clearDisplay();
        display.setTextSize(1);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("Transmission of Score: ");
        display.setCursor(0, 20);
        display.println("Yield Score:");
        display.setTextSize(3); // BIGG
        display.setCursor(0, 35);
        display.println(yieldScore);
        display.display();
        
        bool currentBit = txData[txBitIndex];
        
        //responive animation (as needed) in dot matrix of the currenlty transfering bits
        if (currentBit) {
          lc.setRow(0, 0, B00011000);
          lc.setRow(0, 1, B00011000);
          lc.setRow(0, 2, B00011000); 
          lc.setRow(0, 3, B00011000); 
          lc.setRow(0, 4, B00011000);
          lc.setRow(0, 5, B00011000);
          lc.setRow(0, 6, B00011000);
          lc.setRow(0, 7, B00011000); 
        } else {
          lc.setRow(0, 0, B00111100);
          lc.setRow(0, 1, B01100110);
          lc.setRow(0, 2, B11000011);
          lc.setRow(0, 3, B11000011);
          lc.setRow(0, 4, B11000011);
          lc.setRow(0, 5, B11000011);
          lc.setRow(0, 6, B01100110);
          lc.setRow(0, 7, B00111100);
        }
        break;
      }
    } 

    // back to menu
    if (systemState != lastState) {
      if (systemState == STATE_MENU) {
        lc.clearDisplay(0);
        for (int i = 0; i < 4; i++) digitalWrite(STEPPER_PINS[i], LOW);
      }
      lastState = systemState;
    }
  } 
}