#include "settings.h"
#include "SPI.h"
#include "Adafruit_FT6206.h"
#include "Bitmap.h"
#include <Fonts/FreeSerif9pt7b.h>
#include <Fonts/FreeSerifBold12pt7b.h>
#include <Fonts/FreeSerif18pt7b.h>
#include <AccelStepper.h>

// Initialize stepper motors using AccelStepper
AccelStepper drillRotation(AccelStepper::FULL4WIRE, 2, 3, 4, 5);
AccelStepper depthStepper(AccelStepper::FULL4WIRE, 6, 7, 8, 9);

// Initialize display and touchscreen objects
Adafruit_FT6206 touch = Adafruit_FT6206();

// Define screen states
enum ScreenState { HOME_SCREEN, MANUAL_SCREEN, FIXED_SCREEN, BOTH_CORTICES_SCREEN, SINGLE_Cortex_SCREEN, SECOND_Cortex_SCREEN };

ScreenState currentScreen = HOME_SCREEN;
bool homeButtonEnabled = false, drillOn = false;
unsigned long previousMillis = 0;
float numericalValue = 0.0;


void setup() {
  Serial.begin(9600); // Initialize serial communication at 9600 baud
  tft.begin();
  pinMode(A1, INPUT); // Set the potentiometer pin as input
  pinMode(A2, INPUT); // Set the potentiometer pin as input
  tft.setRotation(3); // Set display rotation
  tft.fillScreen(EGGSHELL_COLOR); // Fill background color

  // Display initial title
  tft.setCursor(50, 20);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Delman Drill");
  tft.setFont(&FreeSerif18pt7b);

  // Draw initial buttons
  drawButton(BUTTON1_X, BUTTON1_Y, "1.Manual");
  drawSqButton(BUTTON2_X, BUTTON2_Y, "2.Both", "Cortices");
  drawSqButton(BUTTON3_X, BUTTON3_Y, "3.Single", "Cortex");
  drawSqButton(BUTTON4_X, BUTTON4_Y, "4.Fixed", "Depth");
  drawSqButton(BUTTON5_X, BUTTON5_Y, "5.Second", "Cortex Find");

  // Initialize touchscreen
  if (!touch.begin()) {
    tft.setFont(&FreeSerif9pt7b);
    tft.println("Unable to start touchscreen.");
    while (1);
  }

  // Initialize the stepper motors
  drillRotation.setMaxSpeed(2000); // Set max speed to a high value
  drillRotation.setAcceleration(500); // Set a reasonable acceleration
  depthStepper.setMaxSpeed(200); // Set max speed for depth stepper
  depthStepper.setAcceleration(20); // Set acceleration for depth stepper
}

void loop() {
  if (drillOn) {
    drillRotation.runSpeed();
  }

  if (currentScreen == MANUAL_SCREEN) {
    updateManualScreen();
  }

  if (touch.touched()) {
    handleTouchInput();
  }

  
  if (analogRead(A1) != 0) { // Button was just pressed
    if (currentScreen == FIXED_SCREEN) {
      fixedDepthRun();
    }
  }
}

// Draw a button with specified parameters
void drawButton(int x, int y, String label) {
  tft.fillRoundRect(x + 3, y - 1, BUTTON_WIDTH, BUTTON_HEIGHT, 8, BUTTON_BORDER_COLOR);
  tft.fillRoundRect(x, y, BUTTON_WIDTH, BUTTON_HEIGHT, 8, BUTTON_FILL_COLOR);
  tft.drawRoundRect(x, y, BUTTON_WIDTH, BUTTON_HEIGHT, 8, BUTTON_BORDER_COLOR);
  tft.setCursor(x + 90, y + 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerifBold12pt7b);
  tft.println(label);
}

void drawSqButton(int x, int y, String label, String line2) {
  tft.fillRoundRect(x + 3, y - 1, SQ_BUTTON_WIDTH, SQ_BUTTON_HEIGHT, 8, BUTTON_BORDER_COLOR);
  tft.fillRoundRect(x, y, SQ_BUTTON_WIDTH, SQ_BUTTON_HEIGHT, 8, BUTTON_FILL_COLOR);
  tft.drawRoundRect(x, y, SQ_BUTTON_WIDTH, SQ_BUTTON_HEIGHT, 8, BUTTON_BORDER_COLOR);
  int pixLength = (line2.length()-5) * 5;
  tft.setCursor(x + 30, y + 30);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerifBold12pt7b);
  tft.println(label);
  tft.setCursor(x+45-pixLength, y + 60);
  tft.println(line2);
}

// Handle touch input
void handleTouchInput() {
  TS_Point p = touch.getPoint();
  // Adjust touch coordinates for the rotation of the display
  int16_t x = tft.width() - map(p.y, 0, 240, 0, tft.height());
  int16_t y = map(p.x, 0, 320, 0, tft.width());

  switch (currentScreen) {
    case HOME_SCREEN:
      // Handle button presses on home screen
      if (x >= BUTTON1_X && x <= BUTTON1_X + BUTTON_WIDTH && y >= BUTTON1_Y && y <= BUTTON1_Y + BUTTON_HEIGHT) changeScreen(MANUAL_SCREEN);
      else if (x >= BUTTON2_X && x <= BUTTON2_X + SQ_BUTTON_WIDTH && y >= BUTTON2_Y && y <= BUTTON2_Y + SQ_BUTTON_HEIGHT) changeScreen(BOTH_CORTICES_SCREEN);
      else if (x >= BUTTON3_X && x <= BUTTON3_X + SQ_BUTTON_WIDTH && y >= BUTTON3_Y && y <= BUTTON3_Y + SQ_BUTTON_HEIGHT) changeScreen(SINGLE_Cortex_SCREEN);
      else if (x >= BUTTON4_X && x <= BUTTON4_X + SQ_BUTTON_WIDTH && y >= BUTTON4_Y && y <= BUTTON4_Y + SQ_BUTTON_HEIGHT) changeScreen(FIXED_SCREEN);
      else if (x >= BUTTON5_X && x <= BUTTON5_X + SQ_BUTTON_WIDTH && y >= BUTTON5_Y && y <= BUTTON5_Y + SQ_BUTTON_HEIGHT) changeScreen(SECOND_Cortex_SCREEN);
      break;
    case FIXED_SCREEN:
      // Handle home button press on fixed depth screen
      if (homeButtonEnabled && x >= HOME_BUTTON_X && x <= HOME_BUTTON_X + home_width && y >= HOME_BUTTON_Y && y <= HOME_BUTTON_Y + home_height) changeScreen(HOME_SCREEN);
      else handleArrowTouchEvent(x, y); // Handle arrow touch event on fixed depth screen
      break;
    case MANUAL_SCREEN:
    case BOTH_CORTICES_SCREEN:
    case SINGLE_Cortex_SCREEN:
    case SECOND_Cortex_SCREEN:
      // Handle home button press on other screens
      if (homeButtonEnabled && x >= HOME_BUTTON_X && x <= HOME_BUTTON_X + home_width && y >= HOME_BUTTON_Y && y <= HOME_BUTTON_Y + home_height) changeScreen(HOME_SCREEN);
      break;
  }
}

void updateManualScreen() {
  int potValue = analogRead(A1); // Read potentiometer value from pin A1
  int potValue2 = analogRead(A0);
  if (potValue != 0){
    int speed = map(potValue, 0, 1023, 0, 2000); // Map the potentiometer value to a higher speed range
    drillRotation.setSpeed(speed); // Set the stepper speed
    drillRotation.runSpeed(); // Ensure the motor runs at the set speed
  }
  else if (potValue2 != 0){
    int speed = map(potValue2, 0, 1023, 0, -2000); // Map the potentiometer value to a higher speed range
    drillRotation.setSpeed(speed); // Set the stepper speed
    drillRotation.runSpeed(); // Ensure the motor runs at the set speed
  }
}

void updateDepth() {
  tft.fillRect(58, 230, 105, -40, EGGSHELL_COLOR);
  float currentVal = depthStepper.currentPosition() / (1000 / maxDepth);
  tft.setFont(&FreeSerif18pt7b);
  tft.setCursor(99, 220);
  tft.print(currentVal, 1);
}

void fixedDepthRun() {
  float depth = numericalValue;
  unsigned long spinInterval = 2000; // 2 seconds before depthStepper starts moving
  unsigned long startMillis = millis();
  unsigned long lastUpdate = millis(); // Track the last time updateDepth was called
  const unsigned long updateInterval = 500; // 0.5 seconds for update intervals

  // Start the drill rotation motor at full speed
  drillRotation.setMaxSpeed(1000);
  drillRotation.setAcceleration(200); // Adjusted acceleration for smoother operation
  drillRotation.setSpeed(1000);

  // Run the drill rotation motor for the specified spin interval
  while (millis() - startMillis < spinInterval) {
    drillRotation.runSpeed();
  }

  // Initialize depth stepper motor parameters for smooth motion
  depthStepper.setMaxSpeed(1000); // Adjusted max speed for smoother operation
  int stepsForDepth = depth * (1000 / maxDepth); // Convert depth to steps

  // Move the depth stepper motor to the calculated position with smooth acceleration
  depthStepper.moveTo(stepsForDepth);

  // Keep the drill rotation motor spinning while the depth stepper motor moves
  drillOn = true;

  // Loop to run both motors smoothly and update the display periodically
  while (depthStepper.distanceToGo() != 0) {
    drillRotation.runSpeed();
    depthStepper.run();
  }

  // Ramp down the speed of the depth stepper motor gradually to avoid abrupt stops
  depthStepper.setAcceleration(20); // Lower acceleration for smoother deceleration
  updateDepth();
  depthStepper.moveTo(0); // Move back to the starting position
  while (depthStepper.distanceToGo() != 0) {
    drillRotation.runSpeed();
    depthStepper.run();
  }

  // Ensure the stepper motor is exactly at position 0
  if (depthStepper.currentPosition() != 0) {
    depthStepper.moveTo(0);
    while (depthStepper.distanceToGo() != 0) {
      depthStepper.run();
    }
  }

  // Stop the drill rotation motor
  drillOn = false;

  // Reset acceleration to default value
  depthStepper.setAcceleration(10);
}


// Define the missing changeScreen function
void changeScreen(ScreenState newScreen) {
  delay(DEBOUNCE_DELAY); // Short delay to avoid residual touch
  switch (newScreen) {
    case HOME_SCREEN: home(); break;
    case MANUAL_SCREEN: manual(); break;
    case FIXED_SCREEN: fixedDepth(); break;
    case BOTH_CORTICES_SCREEN: bothCortices(); break;
    case SINGLE_Cortex_SCREEN: singleCortex(); break;
    case SECOND_Cortex_SCREEN: secondCortex(); break;
  }
  currentScreen = newScreen;
  homeButtonEnabled = (newScreen != HOME_SCREEN);
}

// Home screen setup
void home() {
  tft.fillScreen(EGGSHELL_COLOR);
  numericalValue = 0;
  tft.setCursor(50, 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Delman Drill");
  tft.setFont(&FreeSerif18pt7b);
  drawButton(BUTTON1_X, BUTTON1_Y, "1.Manual");
  drawSqButton(BUTTON2_X, BUTTON2_Y, "2.Both", "Cortices");
  drawSqButton(BUTTON3_X, BUTTON3_Y, "3.Single", "Cortex");
  drawSqButton(BUTTON4_X, BUTTON4_Y, "4.Fixed", "Depth");
  drawSqButton(BUTTON5_X, BUTTON5_Y, "5.Second", "Cortex Find");

  depthStepper.setAcceleration(10);
  depthStepper.moveTo(0);
  while (depthStepper.distanceToGo() != 0) {
    depthStepper.run();
  }

  // Ensure the stepper motor is exactly at position 0
  if (depthStepper.currentPosition() != 0) {
    depthStepper.moveTo(0);
    while (depthStepper.distanceToGo() != 0) {
      depthStepper.run();
    }
  }
}


// Manual screen setup
void manual() {
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(50, 125);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Setting Up...");
  tft.setFont(&FreeSerif9pt7b);
  tft.setCursor(100, 155);
  tft.println("Please wait");
  depthStepper.moveTo(1001);
  while(depthStepper.distanceToGo() != 0) {
    depthStepper.run();
  }
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(60, 125);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Manual Mode");
  tft.setFont(&FreeSerif18pt7b);
  drawHomeButton();
}

// Fixed depth screen setup
void fixedDepth() {
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(60, 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Fixed Depth");
  tft.setFont(&FreeSerif18pt7b);
  drawArrows(); // Draw up and down arrows
  tft.setCursor(130, 180);
  tft.setFont(&FreeSerif9pt7b);
  tft.setTextColor(GRID_COLOR);
  tft.println("Depth");
  // mm setting
  tft.setCursor(113, 130);
  tft.setFont(&FreeSerif18pt7b);
  tft.setTextColor(TEXT_COLOR);
  tft.print(numericalValue, 1);
  tft.setCursor(165, 130);
  tft.setFont(&FreeSerif9pt7b);
  tft.print("mm");
  tft.setFont();
  tft.setCursor(43, 121);
  tft.print("Coarse");
  tft.setCursor(218, 121);
  tft.print("Fine");

  float potValue = depthStepper.currentPosition() / 100;
  char num[5];
  dtostrf(potValue, 5, 1, num); // Convert value to string
  String val = String(num);
  val.trim();
  int pixelLength = (val.length() - 1) * 18;
  tft.fillRect(58, 230, 105, -40, EGGSHELL_COLOR); // Clear previous value
  tft.setFont(&FreeSerif18pt7b);
  tft.setCursor(145 - pixelLength, 220);
  tft.print(val); // Display new value
  tft.setCursor(165, 220);
  tft.setFont(&FreeSerif9pt7b);
  tft.print(" mm");
  drawHomeButton();
}

// Both cortices screen setup
void bothCortices() {
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(35, 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Both Cortices");
  tft.setFont(&FreeSerif18pt7b);
  drawGraph();
}

// Single Cortex screen setup
void singleCortex() {
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(30, 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Single Cortex");
  tft.setFont(&FreeSerif18pt7b);
  drawGraph();
}

// Second Cortex screen setup
void secondCortex() {
  tft.fillScreen(EGGSHELL_COLOR);
  tft.setCursor(15, 25);
  tft.setTextColor(TEXT_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.println("Second Cortex Find");
  tft.setFont(&FreeSerif18pt7b);
  drawGraph();
  homeButtonEnabled = true; // Enable the home button after drawing the graph
}

// Draw arrows for fixed depth screen
void drawArrows() {
  tft.fillTriangle(ARROW_UP_X, ARROW_UP_Y + ARROW_SIZE, ARROW_UP_X + ARROW_SIZE / 2, ARROW_UP_Y, ARROW_UP_X + ARROW_SIZE, ARROW_UP_Y + ARROW_SIZE, GRID_COLOR);
  tft.fillTriangle(ARROW_DOWN_X, ARROW_DOWN_Y, ARROW_DOWN_X + ARROW_SIZE / 2, ARROW_DOWN_Y + ARROW_SIZE, ARROW_DOWN_X + ARROW_SIZE, ARROW_DOWN_Y, GRID_COLOR);
  tft.fillTriangle(CARROW_UP_X, CARROW_UP_Y + CARROW_SIZE, CARROW_UP_X + CARROW_SIZE / 2, CARROW_UP_Y, CARROW_UP_X + CARROW_SIZE, CARROW_UP_Y + CARROW_SIZE, AXIS_COLOR);
  tft.fillTriangle(CARROW_DOWN_X, CARROW_DOWN_Y, CARROW_DOWN_X + CARROW_SIZE / 2, CARROW_DOWN_Y + CARROW_SIZE, CARROW_DOWN_X + ARROW_SIZE, CARROW_DOWN_Y, AXIS_COLOR);
  tft.fillRoundRect(START_X-8, START_Y-7, START_SIZE + 14, START_SIZE + 14, 6, ILI9341_BLACK);
  tft.fillRoundRect(START_X-6, START_Y-5, START_SIZE + 10, START_SIZE + 10, 6, ILI9341_WHITE);
  tft.fillTriangle(START_X, START_Y, START_X, START_Y + START_SIZE, START_X + START_SIZE, START_Y + START_SIZE/2, ILI9341_GREEN);
}

// Handle touch events on arrows for fixed depth adjustment
void handleArrowTouchEvent(int touchX, int touchY) {
  if (touchX >= ARROW_UP_X && touchX <= ARROW_UP_X + ARROW_SIZE && touchY >= ARROW_UP_Y && touchY <= ARROW_UP_Y + ARROW_SIZE) {
    if(numericalValue < maxDepth) numericalValue += 0.1;
    updateDisplay(); // Update the display with the new numerical value
  }
  if (touchX >= ARROW_DOWN_X && touchX <= ARROW_DOWN_X + ARROW_SIZE && touchY >= ARROW_DOWN_Y && touchY <= ARROW_DOWN_Y + ARROW_SIZE) {
    if (numericalValue >= 0.1) numericalValue -= 0.1;
    updateDisplay(); // Update the display with the new numerical value
  }
  if (touchX >= CARROW_UP_X && touchX <= CARROW_UP_X + CARROW_SIZE && touchY >= CARROW_UP_Y && touchY <= CARROW_UP_Y + CARROW_SIZE) {
    if (int(numericalValue) < maxDepth) numericalValue++;
    updateDisplay(); // Update the display with the new numerical value
  }
  if (touchX >= CARROW_DOWN_X && touchX <= CARROW_DOWN_X + CARROW_SIZE && touchY >= CARROW_DOWN_Y && touchY <= CARROW_DOWN_Y + CARROW_SIZE) {
    if (numericalValue >= 1) numericalValue--;
    updateDisplay(); // Update the display with the new numerical value
  }
  if (touchX >= START_X && touchX <= START_X + START_SIZE && touchY >= START_Y && touchY <= START_Y + START_SIZE) fixedDepthRun();
}

// Draw graph for cortices screens
void drawGraph() {
  drawHomeButton(); // Draw the home button
  // Draw the graph area
  tft.fillRect(5, 35, 294, 145, GRAPH_BG_COLOR);
  // Draw the y-axis
  int ySteps = MAX_Y_AXIS_VALUE / Y_AXIS_INCREMENT;
  for (int i = 0; i <= ySteps; i++) {
    int y = GRAPH_Y + GRAPH_HEIGHT - (i * (GRAPH_HEIGHT / ySteps));
    tft.drawLine(GRAPH_X, y, GRAPH_X - 2, y, GRAPH_LINE_COLOR); // Draw horizontal grid lines
    drawMultiDigitNumberBitmap(GRAPH_X - num_width * 3 + 4, y - num_height / 2, i * Y_AXIS_INCREMENT); // Right align numbers
  }
  // Draw the x-axis
  int xSteps = MAX_X_AXIS_VALUE / X_AXIS_INCREMENT;
  for (int i = 0; i <= xSteps; i++) {
    int x = GRAPH_X + (i * (GRAPH_WIDTH / xSteps));
    tft.drawLine(x, GRAPH_Y + GRAPH_HEIGHT, x, GRAPH_Y + GRAPH_HEIGHT + 2, GRAPH_LINE_COLOR); // Draw vertical grid lines
    drawMultiDigitNumberBitmap(x, GRAPH_Y + GRAPH_HEIGHT + 4, i * X_AXIS_INCREMENT); // Right align numbers
  }
  // Draw the axes
  tft.drawLine(GRAPH_X, GRAPH_Y, GRAPH_X, GRAPH_Y + GRAPH_HEIGHT, GRAPH_AXIS_COLOR); // y-axis
  tft.drawLine(GRAPH_X, GRAPH_Y + GRAPH_HEIGHT, GRAPH_X + GRAPH_WIDTH, GRAPH_Y + GRAPH_HEIGHT, GRAPH_AXIS_COLOR); // x-axis
  // Draw axis labels
  tft.setRotation(4);
  tft.setCursor(105, 15);
  tft.setFont();
  tft.setTextColor(TEXT_COLOR);
  tft.print("Current (mA)"); // y-axis label
  tft.setRotation(3);
  tft.setCursor(130, 165);
  tft.print("Time (s)"); // x-axis label
  drawNumberLabels(); // Draw additional labels below the graph
}

// Draw additional labels below the graph
void drawNumberLabels() {
  tft.setFont(&FreeSerif9pt7b);
  tft.setTextColor(TEXT_COLOR);
  // Position
  tft.setCursor(20, 195);
  tft.print(POSITION_LABEL);
  tft.setCursor(86, 195);
  tft.print("0"); // Placeholder value
  // Time
  tft.setCursor(120, 195);
  tft.print(TIME_LABEL);
  tft.setCursor(165, 195);
  tft.print("0"); // Placeholder value
  // Current
  tft.setCursor(200, 195);
  tft.print(CURRENT_LABEL);
  tft.setCursor(262, 195);
  tft.print("0"); // Placeholder value
}

// Draw the home button
void drawHomeButton() {
  tft.fillRect(HOME_BUTTON_X + 6, HOME_BUTTON_Y + 6, home_width - 12, home_height - 12, HOME_BUTTON_FILL_COLOR);
  tft.drawBitmap(HOME_BUTTON_X, HOME_BUTTON_Y, home_bits, home_width, home_height, HOME_BUTTON_BORDER_COLOR);
}

// Update display with numerical value
void updateDisplay() {
  delay(UPDATE_DELAY);
  int val = (String(numericalValue).length() - 1) * 18;
  tft.fillRect(78, 150, 82, -50, EGGSHELL_COLOR);
  tft.setFont(&FreeSerif18pt7b);
  tft.setCursor(166 - val, 130);
  tft.setTextColor(TEXT_COLOR);
  tft.print(numericalValue, 1);
  delay(150);
}

// Draw number bitmap for a single digit
void drawNumberBitmap(int x, int y, int number) {
  const uint8_t *bitmap;
  switch (number) {
    case 0: bitmap = zero_bits; break;
    case 1: bitmap = one_bits; break;
    case 2: bitmap = two_bits; break;
    case 3: bitmap = three_bits; break;
    case 4: bitmap = four_bits; break;
    case 5: bitmap = five_bits; break;
    case 6: bitmap = six_bits; break;
    case 7: bitmap = seven_bits; break;
    case 8: bitmap = eight_bits; break;
    case 9: bitmap = nine_bits; break;
    default: return; // invalid number
  }
  tft.drawBitmap(x, y, bitmap, num_width, num_height, TEXT_COLOR);
}

// Draw multi-digit number bitmap
void drawMultiDigitNumberBitmap(int x, int y, int number) {
  char str[4]; // Buffer to hold the number string, 3 digits + null terminator
  itoa(number, str, 10); // Convert the number to a string
  int len = strlen(str); // Get the length of the string

  // Draw each digit, from right to left
  for (int i = len - 1; i >= 0; i--) {
    drawNumberBitmap(x - (len - 1 - i) * num_width, y, str[i] - '0');
  }
}