#include "HX711.h"

// Pin definitions
const byte LOADCELL_DOUT_PIN = 34;
const byte LOADCELL_SCK_PIN = 23;
const byte FW_RELAY_PIN = 32;
const byte BW_RELAY_PIN = 33;
const byte BUTTON_PIN=4;

// Machine states definition
enum states {   
  NONE,RESET, 
  EXCEPTION,
  MANUAL,
  COMPRESS,
  PULL,
  PAUSE,
  COMPLETE,
};
// Machine substates definition
enum substates {
  THRUST,RETRACT,HOLD,S_NONE,S_DONE,S_PAUSE
};

// Instanciate the scale object
HX711 scale;

// Calibration constants
const float SCALE_CALIBRATION_FACTOR = 400.0;
const float INITIAL_TARE_THRESHOLD = 0.2;
const int NUM_READINGS_AVG = 2;

// Status checks and flags
bool calibIsOk=true;
bool HX711IsOK=true;
String currentState="";
byte holdType=0;
bool pauseCycle=false;
bool skipSerialRead=false;
unsigned long previousPingMillis = 0;
const long PING_INTERVAL = 1000;

// Variables for target cycles and force
int targetCycles = 0;
float targetForce = 0.0;
int targetHold=0;
int defaultHold=200;
int completedCycles = 0;
int forceSign=0;


// Global variables to store the prior and current states
states prior_state, state; 
substates prior_substate, substate,paused_substate;

void reset() {
  if (state != prior_state) {
  // Functions performed when entering state
  prior_state = state;
  currentState="Idle";
  }
  // Perform state tasks
  if (Serial.available() > 0) {
    String command = Serial.readStringUntil('\n');
    command.trim();
    char com_char=command[0];
    char cmd;
    int arg1,arg3;
    float arg2;
    switch(com_char){
      case 'r':
        ESP.restart();
      case 'm':
        state=MANUAL;
      break;
      case 'c':
        sscanf(command.c_str(), "%c,%d,%f,%d", &cmd, &arg1, &arg2,&arg3);
        state=COMPRESS;
        targetCycles=arg1;
        targetForce=arg2;
        targetHold=arg3;
      break;
      case 'p':
        sscanf(command.c_str(), "%c,%d,%f,%d", &cmd, &arg1, &arg2,&arg3);
        state=PULL;
        targetCycles=arg1;
        targetForce=arg2;
        targetHold=arg3;
      break;
    }
  }

   // Check for state transitions
  if (state != prior_state) {
// Functions performed when leaving state
  }
}
void calibrate() {
  // Perform state tasks
  scale.tare();
  scale.set_scale(SCALE_CALIBRATION_FACTOR);
  delay(1000);
  float tareResult = scale.get_units(10);
  if (abs(tareResult) > INITIAL_TARE_THRESHOLD) {
    calibIsOk=false;}
}
void sensorCheck() {
  scale.begin(LOADCELL_DOUT_PIN, LOADCELL_SCK_PIN);
  // Ensure the scale has stabilized
  delay(2000);
  // Perform state tasks
  bool allZero = true;
  // Get 10 readings and check if all are zero
  for (int i = 0; i < 10; i++) {
    if (scale.is_ready()) {
      long weight = scale.get_units(2);  // Get a single reading
      // If any reading is not zero, set allZero to false and break out of the loop
      if (weight != 0) {
        allZero = false;
        break;  // No need to continue checking if we already found a non-zero reading
      }
    }
  }
  if (allZero){
    HX711IsOK=false;
  }
}

void exception(){
  // If we are entering the state, do initialization stuff
  if (state != prior_state) {         
    prior_state = state;
    currentState="Exception";
  }
  //state tasks
  if (Serial.available() > 0) {
    String command = Serial.readStringUntil('\n');
    command.trim();
      if (command == "r") {
        ESP.restart();
      }
  }
}
void manual() {
  float currentForce=0.0;
  if (state != prior_state) {// If we are entering the state, do initialization stuff
    prior_state = state;
    currentState="Manual mode";
  }
  
  // Perform state tasks
  currentForce = scale.get_units(NUM_READINGS_AVG);
  Serial.print("M,");
  Serial.println(currentForce);

    if (Serial.available() > 0) {
    String command = Serial.readStringUntil('\n');
    command.trim();
      if (command.startsWith("F") ||command.startsWith("B")) {
        String valueStr =command.substring(2,command.length());
        uint16_t value_ms = (uint16_t)valueStr.toInt(); // Convert to uint16_t
        move_till_time(command.charAt(0),value_ms);
      }
      if (command=="s") {
        state=RESET;
      }
      
  }
  // Check for state transitions
  if (state != prior_state) {// If we are leaving the state, do clean up stuff
  }
}

void s_pause(){
//if we are entering state, perform initialization tasks
if (substate != prior_substate){
skipSerialRead=true;
digitalWrite(FW_RELAY_PIN, 0);
digitalWrite(BW_RELAY_PIN, 0);
currentState="Pause";
paused_substate=prior_substate;
prior_substate = substate;
}

// Perform state tasks
  if (Serial.available() > 0) {
  String command = Serial.readStringUntil('\n');
  command.trim();
    //resume cyclic test
    if (command=="a") {
      substate=paused_substate;
    }
    //save data and end test
    if (command=="b") {
      substate=paused_substate;
      targetCycles=completedCycles;
    }
    
}
// If we are leaving the state, do clean up stuff
if (substate != prior_substate) {
currentState="Compress mode";
skipSerialRead=false;
}
}


void compress(int nCycles,float tForce, int hold_MS) {
  if (state != prior_state) {// If we are entering the state, do initialization stuff
    prior_state = state;
    currentState="Compress mode";
    prior_substate=S_NONE;
    substate=THRUST;
    forceSign=1;
  }
    // Perform state tasks
  if (Serial.available() > 0 && !skipSerialRead) {
    String command = Serial.readStringUntil('\n');
    command.trim();
      if (command=="w") {
        substate=S_PAUSE;
      }
  }
  if(completedCycles<nCycles){
    switch(substate){
      case S_PAUSE:
        s_pause();
        break;
      case THRUST:
        thrust(tForce);
        break;
      case HOLD:
        if(holdType==2){
          hold(defaultHold,holdType);
        }
        else{
          hold(hold_MS,holdType);
        }
        break;
      case RETRACT:
        retract(0.2*tForce);
        break;
    }
  }
  else if(completedCycles>=targetCycles){
  state=RESET;
  Serial.print("R,0,");
  Serial.println(completedCycles);
  }
  // Check for state transitions
// If we are leaving the state, do clean up stuff
  if (state != prior_state) {
    completedCycles=0;
    skipSerialRead=false;
    Serial.println("MSG,Test complete");
    forceSign=0;
  }
}

void pull(int nCycles,float tForce, int hold_MS) {
  if (state != prior_state) {// If we are entering the state, do initialization stuff
    prior_state = state;
    currentState="Pull mode";
    prior_substate=S_NONE;
    substate=RETRACT;
    forceSign=-1;
  }
    // Perform state tasks
  if (Serial.available() > 0 && !skipSerialRead) {
    String command = Serial.readStringUntil('\n');
    command.trim();
      if (command=="w") {
        substate=S_PAUSE;
      }
  }
  if(completedCycles<nCycles){
    switch(substate){
      case S_PAUSE:
        s_pause();
        break;
      case THRUST:
        thrust(-0.2*tForce);
        break;
      case HOLD:
        if(holdType==2){
          hold(hold_MS,holdType);
        }
        else{
          hold(defaultHold,holdType);
        }
        break;
      case RETRACT:
        retract(-tForce);
        break;
    }
  }
  else if(completedCycles>=targetCycles){
  state=RESET;
  Serial.print("R,0,");
  Serial.println(completedCycles);
  }
  // Check for state transitions
// If we are leaving the state, do clean up stuff
  if (state != prior_state) {
    completedCycles=0;
    skipSerialRead=false;
    Serial.println("MSG,Test complete");
    forceSign=0;
  }
}

void move_till_time(char dir,uint16_t step_ms){
  unsigned long start_time=0;
  unsigned long current_time=0;
  if(dir=='F'){
    digitalWrite(FW_RELAY_PIN, 1);
    start_time=millis();
    current_time=start_time;
    while(current_time-start_time<step_ms){
      current_time=millis();
    }
    digitalWrite(FW_RELAY_PIN, 0);
  }
  if(dir=='B'){
    digitalWrite(BW_RELAY_PIN, 1);
    start_time=millis();
    current_time=start_time;
    while(current_time-start_time<step_ms){
      current_time=millis();
    }
    digitalWrite(BW_RELAY_PIN, 0);
  }
}

void thrust(float tForce) {
  float currentForce=0.0;
  if (substate != prior_substate) {// If we are entering the state, do initialization stuff
    prior_substate = substate;
    digitalWrite(FW_RELAY_PIN, 1);
  }
  // Perform state tasks
  currentForce =scale.get_units(NUM_READINGS_AVG);
  Serial.print("R,");
  Serial.print(currentForce*forceSign);
  Serial.print(",");
  Serial.println(completedCycles);
  // Check for state transitions
  if (currentForce>=tForce) {
    substate=HOLD;
    holdType=1;
  }
  if (substate != prior_substate) {// If we are leaving the state, do clean up stuff
    digitalWrite(FW_RELAY_PIN, 0);
    if(forceSign==-1){completedCycles++;} 
  }
}
void retract(float tForce) {
  float currentForce=0.0;
  if (substate != prior_substate) {// If we are entering the state, do initialization stuff
    prior_substate = substate;
    digitalWrite(BW_RELAY_PIN, 1);
  }
  // Perform state tasks
  currentForce = scale.get_units(NUM_READINGS_AVG);
  Serial.print("R,");
  Serial.print(currentForce*forceSign);
  Serial.print(",");
  Serial.println(completedCycles);
  if (currentForce<= tForce) {
    substate=HOLD;
    holdType=2;
  }
  // Check for state transitions
  if (substate != prior_substate) {// If we are leaving the state, do clean up stuff
    digitalWrite(BW_RELAY_PIN, 0);
    if(forceSign==1){completedCycles++;} 
  }
}

void hold(int step_ms,byte holdType){
  static unsigned long start_time = 0; // Use static so the value persists between calls
  unsigned long current_time = 0; // Get the current time
  if (substate != prior_substate) {// If we are entering the state, do initialization stuff
  prior_substate = substate;
  start_time = millis();
  }
  current_time=millis();
  if(current_time-start_time>=step_ms){
      if (holdType==1){
        substate=RETRACT;
      }
      else{
        substate=THRUST;
      }
    }
  if (substate != prior_substate) {// If we are leaving the state, do clean up stuff
  }
}

void IRAM_ATTR buttonPressHandler() {
pauseCycle=true;
//digitalWrite(FW_RELAY_PIN, 0);
//digitalWrite(BW_RELAY_PIN, 0);
}
void setup() {
  // Initialize relay pins
  pinMode(FW_RELAY_PIN, OUTPUT);
  pinMode(BW_RELAY_PIN, OUTPUT);
  digitalWrite(FW_RELAY_PIN, 0);
  digitalWrite(BW_RELAY_PIN, 0);
  prior_state = NONE;
    // Initialize serial communication
  Serial.begin(115200);
  pinMode(BUTTON_PIN, INPUT_PULLDOWN); // Enable internal pull-up resistor
  attachInterrupt(digitalPinToInterrupt(BUTTON_PIN), buttonPressHandler, FALLING);
  sensorCheck();
  calibrate();
  if(calibIsOk && HX711IsOK){
      state = RESET;
  }
  else{
    state=EXCEPTION;
  }
}

void loop() {
  switch (state) {
    case RESET:
      reset();
      break;
    case EXCEPTION:
      exception();
      break;
    case MANUAL:
      manual();
      break;
    case COMPRESS:
      compress(targetCycles,targetForce,targetHold);
      break;
    case PULL:
      pull(targetCycles,targetForce,targetHold);
      break;
  }
  sendPing();
}

void sendPing() {
  unsigned long currentPingMillis = millis();  // Get the current time

  // Check if 10 seconds have passed since the last "PING"
  if (currentPingMillis - previousPingMillis >= PING_INTERVAL) {
    // Save the current time
    previousPingMillis = currentPingMillis;
    
    // Send the "PING" message to the serial port
    Serial.print("PING,");
    Serial.println(currentState);
  }
}