#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <Keypad.h>
#include <AccelStepper.h>

// Define stepper motor connections and motor interface type
#define STEP_PIN 11
#define DIR_PIN 10
#define ENABLE_PIN 12

#define STEP_PER_REV 200
#define MICROSTEP_PER_STEP 1 // Increase microstepping for higher resolution
#define MICROSTEP_PER_REV (STEP_PER_REV * MICROSTEP_PER_STEP)
#define SCREW_PITCH_MM 8
#define GEARBOX 1

// Define syringe diameter
#define SYRINGE_DIAMETER_MM 20

// LCD and Keypad setup
LiquidCrystal_I2C lcd(0x27, 20, 4);

const byte ROWS = 4; // four rows
const byte COLS = 3; // three columns
char keys[ROWS][COLS] = {
  {'1','2','3'},
  {'4','5','6'},
  {'7','8','9'},
  {'*','0','#'}
};
byte rowPins[ROWS] = {2,3,4,6}; // connect to the row pinouts of the keypad
byte colPins[COLS] = {7,8,9}; // connect to the column pinouts of the keypad

Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// Create stepper motor object
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

enum State { STATE_IDLE, GET_INPUT_VOLUME, GET_INPUT_FLOWRATE, CHOOSE, INFUSE, DEFUSE, RETURN, ERROR, WELCOME };
enum Event { INPUT_COMPLETE, CHOOSE_COMPLETE, INFUSE_START, DEFUSE_START, FINISHED, ERROR_OCCURRED, RESET, NO_EVENT, BACK };

State currentState = WELCOME;
State nextState = WELCOME;
State previousState = WELCOME;

float volume = 0;
float flowRate = 0;
char operation = ' ';
String inputBuffer = "";
bool starPressedOnce = false;
unsigned long starPressTime = 0;
unsigned long welcomeStartTime = 0;
const unsigned long doublePressInterval = 1000; // 1 second for double press
const unsigned long welcomeDuration = 5000; // 5 seconds for welcome message

Event checkExternalEvent(void);
void doAction(State currentState);
void exitAction(State currentState);
void entryAction(State nextState);
void stateMachine(Event event);

uint32_t calculateSteps(float volume_cm3, float syringeDiameter_mm);
float SP_calculateSpeed(float flowRate_cm3_min);

void setup() {
  pinMode(ENABLE_PIN, OUTPUT);
  digitalWrite(ENABLE_PIN, LOW); // Enable the driver

  stepper.setMaxSpeed(1000);  // Adjust max speed if needed
  stepper.setAcceleration(500); // Adjust acceleration for smoother motion

  lcd.init();
  lcd.backlight();
  
  Serial.begin(9600); // Initialize serial communication for debugging
  entryAction(WELCOME);
}

void loop() {
  stateMachine(checkExternalEvent());
}

Event checkExternalEvent(void) {
  char key = keypad.getKey();
  if (key) {
    Serial.print("Key pressed: ");
    Serial.println(key);
    
    if (key == '#') { // Confirm input
      starPressedOnce = false; // Reset double press tracking
      switch (currentState) {
        case GET_INPUT_VOLUME:
          volume = inputBuffer.toFloat();
          Serial.print("Volume set to: ");
          Serial.println(volume);
          inputBuffer = "";
          return INPUT_COMPLETE;
        case GET_INPUT_FLOWRATE:
          flowRate = inputBuffer.toFloat();
          Serial.print("Flow rate set to: ");
          Serial.println(flowRate);
          inputBuffer = "";
          return INPUT_COMPLETE;
        case CHOOSE:
          if (inputBuffer.length() > 0) {
            operation = inputBuffer[0];
            Serial.print("Operation chosen: ");
            Serial.println(operation);
            inputBuffer = "";
            return CHOOSE_COMPLETE;
          }
        default:
          break;
      }
    } else if (key == '*') { // Reset input or go back to previous state
      if (starPressedOnce && (millis() - starPressTime < doublePressInterval)) {
        starPressedOnce = false; // Reset double press tracking
        return BACK;
      } else {
        starPressedOnce = true;
        starPressTime = millis();
        inputBuffer = "";
        lcd.setCursor(0, 3);
        lcd.print("                    "); // Clear input line
        Serial.println("Input reset");
      }
    } else {
      starPressedOnce = false; // Reset double press tracking
      inputBuffer += key;
      lcd.setCursor(0, 3);
      lcd.print(inputBuffer);
      Serial.print("Input buffer: ");
      Serial.println(inputBuffer);
    }
  }
  return NO_EVENT;
}

void doAction(State state) {
  uint32_t stepsInfuse = 0;
  uint32_t stepsDefuse = 0;
  float speed_cm_min = SP_calculateSpeed(flowRate);
  float speed_mm_min = speed_cm_min * 10; // Convert cm/min to mm/min

  switch (state) {
    case STATE_IDLE:
      // Idle action
      break;
    case GET_INPUT_VOLUME:
      // Prompting for volume input
      break;
    case GET_INPUT_FLOWRATE:
      // Prompting for flow rate input
      break;
    case CHOOSE:
      // Choosing operation action
      break;
    case INFUSE:
      // Infusing action
      stepsInfuse = calculateSteps(volume, SYRINGE_DIAMETER_MM);
      Serial.print("Steps to infuse: ");
      Serial.println(stepsInfuse);
      stepper.setSpeed(-speed_mm_min); // Set speed for infusing (negative for backward)
      Serial.print("Speed set for infusing: ");
      Serial.println(-speed_mm_min);
      digitalWrite(ENABLE_PIN, LOW); // Ensure motor is enabled
      stepper.moveTo(stepper.currentPosition() - stepsInfuse);
      while (stepper.distanceToGo() != 0) {
        stepper.runSpeedToPosition();
      }
      Serial.println("Infusion complete");
      nextState = STATE_IDLE; // Remain in the idle state after infusing
      break;
    case DEFUSE:
      // Defusing action
      stepsDefuse = calculateSteps(volume, SYRINGE_DIAMETER_MM);
      Serial.print("Steps to defuse: ");
      Serial.println(stepsDefuse);
      stepper.setSpeed(speed_mm_min); // Set speed for defusing (positive for forward)
      Serial.print("Speed set for defusing: ");
      Serial.println(speed_mm_min);
      digitalWrite(ENABLE_PIN, LOW); // Ensure motor is enabled
      stepper.moveTo(stepper.currentPosition() + stepsDefuse);
      while (stepper.distanceToGo() != 0) {
        stepper.runSpeedToPosition();
        // Safety check to stop the motor if it goes beyond the expected volume
        if (stepper.currentPosition() >= stepsDefuse) {
          Serial.println("Defusion complete - Safety stop");
          stepper.stop();
          break;
        }
      }
      Serial.println("Defusion complete");
      nextState = STATE_IDLE; // Stay in the idle state after defusing
      break;
    case RETURN:
      // Returning to initial position
      Serial.println("Returning to initial position");
      digitalWrite(ENABLE_PIN, LOW); // Ensure motor is enabled
      stepper.moveTo(0);
      while (stepper.distanceToGo() != 0) {
        stepper.runSpeedToPosition();
      }
      Serial.println("Return complete");
      nextState = STATE_IDLE;
      break;
    case ERROR:
      // Error action
      Serial.println("Error state encountered");
      break;
    case WELCOME:
      // Welcome message
      lcd.setCursor(0, 0);
      lcd.print("SCHOOL OF BME-HCMIU");
      lcd.setCursor(0, 1);
      lcd.print("MED DESIGN 2B");
      lcd.setCursor(0, 2);
      lcd.print("GROUP 3");
      if (millis() - welcomeStartTime >= welcomeDuration) {
        nextState = STATE_IDLE;
      }
      break;
  }
}

void exitAction(State state) {
  // Placeholder for exit actions
}

void entryAction(State state) {
  lcd.clear();
  switch (state) {
    case STATE_IDLE:
      lcd.setCursor(0, 0);
      lcd.print("Enter volume (ml):");
      nextState = GET_INPUT_VOLUME; // Move to the next state
      break;
    case GET_INPUT_VOLUME:
      lcd.setCursor(0, 0);
      lcd.print("Enter volume (ml):");
      break;
    case GET_INPUT_FLOWRATE:
      lcd.setCursor(0, 0);
      lcd.print("Enter flow rate:");
      lcd.setCursor(0, 1);
      lcd.print("(ml/min)");
      break;
    case CHOOSE:
      lcd.setCursor(0, 0);
      lcd.print("Choose operation:");
      lcd.setCursor(0, 1);
      lcd.print("1 for infuse");
      lcd.setCursor(0, 2);
      lcd.print("2 for defuse");
      break;
    case INFUSE:
      lcd.setCursor(0, 0);
      lcd.print("Infusing...");
      break;
    case DEFUSE:
      lcd.setCursor(0, 0);
      lcd.print("Defusing...");
      break;
    case RETURN:
      lcd.setCursor(0, 0);
      lcd.print("Returning...");
      break;
    case ERROR:
      lcd.setCursor(0, 0);
      lcd.print("Error occurred");
      break;
    case WELCOME:
      lcd.setCursor(0, 0);
      lcd.print("SCHOOL OF BME IU");
      lcd.setCursor(0, 1);
      lcd.print("MED DESIGN 2B");
      lcd.setCursor(0, 2);
      lcd.print("GROUP 3");
      welcomeStartTime = millis(); // Start the timer
      break;
  }
  Serial.print("Entering state: ");
  Serial.println(state);
}

void stateMachine(Event event) {
  if (event != NO_EVENT) {
    Serial.print("Event: ");
    Serial.println(event);
  }

  switch (currentState) {
    case STATE_IDLE:
      if (event == INPUT_COMPLETE) nextState = GET_INPUT_VOLUME;
      break;
    case GET_INPUT_VOLUME:
      if (event == INPUT_COMPLETE) nextState = GET_INPUT_FLOWRATE;
      if (event == BACK) nextState = STATE_IDLE;
      break;
    case GET_INPUT_FLOWRATE:
      if (event == INPUT_COMPLETE) nextState = CHOOSE;
      if (event == BACK) nextState = GET_INPUT_VOLUME;
      break;
    case CHOOSE:
      if (event == CHOOSE_COMPLETE) {
        if (operation == '1') nextState = INFUSE;
        else if (operation == '2') nextState = DEFUSE;
        else nextState = ERROR;
      }
      if (event == BACK) nextState = GET_INPUT_FLOWRATE;
      break;
    case INFUSE:
      nextState = STATE_IDLE;
      break;
    case DEFUSE:
      nextState = STATE_IDLE;
      break;
    case RETURN:
      nextState = STATE_IDLE;
      break;
    case ERROR:
      if (event == RESET) nextState = STATE_IDLE;
      break;
    case WELCOME:
      if (millis() - welcomeStartTime >= welcomeDuration) {
        nextState = STATE_IDLE;
      }
      break;
  }

  if (currentState != nextState) {
    previousState = currentState;
    exitAction(currentState);
    currentState = nextState;
    entryAction(nextState);
  }
  doAction(currentState);
}

uint32_t calculateSteps(float volume_cm3, float syringeDiameter_mm) {
  float radius = syringeDiameter_mm / 20.0; // Convert diameter from mm to cm
  float area_cm2 = 3.14 * radius * radius;
  float length_cm = volume_cm3 / area_cm2;
  return (length_cm * 10.0 / SCREW_PITCH_MM) * MICROSTEP_PER_REV;
}

float SP_calculateSpeed(float flowRate_cm3_min) {
  float area_cm2, radius, speed_cm_min;
  radius = SYRINGE_DIAMETER_MM / 20.0; // Convert diameter from mm to cm
  area_cm2 = 3.14 * radius * radius;
  speed_cm_min = flowRate_cm3_min / area_cm2;
  return speed_cm_min;
}