/*
 * Arduino control code for OS-railway Z70
 * This board uses a parkFun Motor Driver - Dual TB6612FNG (1A) for motor control. 
 * The TB6612FNG needs two direction I/Os as they control each output to the motor. The speed is controlled with a PWM output.
 * 
 * Created 2018-10-08 for the OS-railway project 
 * Olle Sköld (Depronized on Thingiverse.com)
 * 
 * Code comes with no liability or copyright, do whatever you want with it. :)
 */

// The code is prepared for using both body lights and roof lights. IF you don't want to implement
// any one of these, you don't have to remove the code, Setting these outputs without LED's connected
// makes no harm

#define headLights 5    // PWM output for front lights
#define aftLights 6     // PWM output for aft lights
#define headRedLight 9  // PWM output for red roof light forward direction
#define aftRedLight 10  // PWN output for red roof light aft direction

#define button 21     // Input for roof button (labled A3 on the Pro Micro).
#define speedInput 12 // Analog input for speed control (labled A2 on the Pro Micro)

// TB6612FNG Motor board
#define motorPin 3    // PWM output for speed control
#define direction1 7  // First direction output 
#define direction2 8  // Second direction output

// State machine flags used to store present state of operation
#define standstil 0
#define acceleration 1
#define cruise 2
#define deceleration 3
char MOTION_STATE = standstil;

#define forward true
#define backward false

//Motion variables for trajectory generator
int ASPD = 0;                  // Actual SPeeD
int TSPD = 0;                  // Target SPeeD
char ACC = 2;                   // Acceleration time in seconds
bool DIR;                       // Direction, forward = true, backward = false
int LOOP_DELAY = 100;          // 100ms loop delay, check button and update speed 10 times per second
char buttonCounter = 0;         // used to count loops with button pressed to initiate emergency stop.

/****************************************
 * If train has derailed or for any other reason needs to stop fast, hold button for more than 
 * the emergencyStopDelay will bypass the deceleration ramp and set motor PWM to zero
 */
char emergencyStopDelay = 1;    // Emergency stop delay seconds.

/****************************************
 *  If button is held longer than the directionDelay, motor is started in reverse direction
 */
char directionDelay = 2;        // Direction delay in seconds. 


bool buttonPressed = false;

/****************************************
 *  Periferal functions
 */

void setMotor(int level){
  if(level > 255){
    level = 255;
  }
  if(level < 0){
    level = 0;
  }
  analogWrite(motorPin, level);  
}

void setDirection(bool dir){
    digitalWrite(direction1, !dir);
    digitalWrite(direction2, dir);  
}

void setup() {

  Serial.begin(115200);

  pinMode(headLights, OUTPUT);
  pinMode(aftLights, OUTPUT);
  pinMode(headRedLight, OUTPUT);
  pinMode(aftRedLight, OUTPUT);
  pinMode(motorPin, OUTPUT);
  pinMode(direction1, OUTPUT);
  pinMode(direction2, OUTPUT);
  pinMode(speedInput, INPUT);
  pinMode(button, INPUT);

  // Set up initial operation for forward direction
  analogWrite(headLights, 255);
  analogWrite(aftRedLight, 255);

  analogWrite(aftLights, 0);
  analogWrite(headRedLight, 0);
  setMotor(0);
  setDirection(forward);
}

void readRoofButton(){
  int temp;
  temp = analogRead(button);
  if(temp > 512){
    buttonPressed = true;
  } else {
    buttonPressed = false;
  }
}

void waitForRelease(){
    while(buttonPressed){
      readRoofButton();
      buttonCounter++;
      delay(LOOP_DELAY);
    }
}

void loop() {

  Serial.println("made it to the loop");
//check button
  readRoofButton();
   Serial.println("made it to readRoofButton");

//State machine
  switch(MOTION_STATE){
    
    case standstil:
      // In case of standstil, a button pressed means start acceleration. 
      // Set state machine to acceleration and find out how fast we want to go.
      // If button is not pressed, don't do anything 
       Serial.println("made it to case standstil");
           
      if(buttonPressed){   
         Serial.println("button press recognised");     
        //Wait for the button to be released before initiating accelration
        waitForRelease();

        //Choose direction depending on the time the button was pressed
        if(buttonCounter < (directionDelay * (1000/LOOP_DELAY))){
          setDirection(forward);
          analogWrite(headLights, 255);
          analogWrite(aftRedLight, 255);
          analogWrite(aftLights, 0);
          analogWrite(headRedLight, 0);
          Serial.println("forward");
        } else {
          setDirection(backward);
          analogWrite(headLights, 0);
          analogWrite(aftRedLight, 0);
          analogWrite(aftLights, 255);
          analogWrite(headRedLight, 255);
          Serial.println("backward");
        }

        //reset buttonCounter;
        buttonCounter = 0;
        
        MOTION_STATE = acceleration;
        
        //read potentiometer and set target speed
        TSPD = (analogRead(speedInput)/4);
        ASPD = 0;
      }
      break;
        
    case acceleration:
      // In acceleration phase, if button is pressed means we want to stop so in that case, set state machine to deceleration
      // If button is not pressed, continue wiht acceleration
            
      if(buttonPressed){
        // Don't wait for release of the button here, we want to initate deceleration without delay
        MOTION_STATE = deceleration;  
      } else {
        
        // If speed is below target speed, accelerate        
        if(ASPD < TSPD){
            ASPD = ASPD + (255 / (ACC*(1000/LOOP_DELAY)));
        }
        
        // If speed after this acceleration is above or equal to target speed, make sure it's not above and set state machine to cruise
        if(ASPD > TSPD){  
          ASPD = TSPD;
          MOTION_STATE = cruise;
        }

      }
      // Update motor speed
      setMotor(ASPD);
      break;
      
    case cruise:
      // In cruising state, check speed input and button
      if(buttonPressed){
        // don't wait for release here as we want to initiate deceleration without delay
        MOTION_STATE = deceleration;  
      }
      // Read analog speed input and adjust speed
      ASPD = (analogRead(speedInput)/4);
      // Update motor speed
      setMotor(ASPD);
      break;

    case deceleration:

      // If button is pressed and held for the duration of the emergencyStopDelay, set speed to 0 and reset motion parameters.
      if(buttonPressed){
        buttonCounter++;
        if(buttonCounter > (emergencyStopDelay*(1000/LOOP_DELAY))){
          setMotor(0);
          MOTION_STATE = standstil;
          ASPD = 0;
          TSPD = 0;
          waitForRelease();
          break;
        }
      } else {
        buttonCounter = 0;
      }

      // We don't use an else here, that is because we want to continue the deceleration even when the emergency stop function is counting.
      // When an emergency stop is initiated, the state machine is exited with the break; command before it reaches here.

      // If speed is still higher than 0, continue deceleration
      if(ASPD > 0){
        ASPD = ASPD - (255 / (ACC*(1000/LOOP_DELAY)));
      } else {
        //If deceleration is finished, reset parameters and set state machine to standstil
        ASPD = 0;
        TSPD = 0;
        MOTION_STATE = standstil;
      }

      setMotor(ASPD);
      
      break;
  }

  //loop delay
  delay(LOOP_DELAY);
}