#include <SPI.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>


// OLED
#define OLED_RESET      4
#define SCREEN_WIDTH    128
#define SCREEN_HEIGHT   64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);


// Motors
const int USER_EMERGENCY_STOP_PIN = 2;
const int USER_DIRECTION_CHANGE_PIN = 3;
const int RGHT_WHEEL_STEP_PIN = 6;
const int LEFT_WHEEL_STEP_PIN = 7;
const int DIRECTION_PIN = 8;
const int EMERGENCY_STOP_PIN = 9;


// HC-SR04
const int PIN_ECHO = 4;
const int PIN_TRIG = 5;


// Timer 1
const float timer1_total_interval = 100;
float timer1_delay_start = millis();
float timer1_delay_running = false;


// Timer 2
const float timer2_total_interval = 800;
float timer2_delay_start = millis();
float timer2_delay_running = false;


// General
float distance_threshold = 10.0;
float steps_to_distance = 0.015;
int total_steps = 0;
int current_step_pos = 0;
int current_dir = 1;
bool emergency_stopped = false;
bool collision_check = false;
float total_travelled = 0;
float current_pos_mtr = 0;


void setup_oled(){
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3D)) {
    for(;;);
  }
  display.display();
  delay(2000);
  display.clearDisplay();

  display.drawPixel(10, 10, SSD1306_WHITE);
  display.display();
  delay(2000);
}


void setup_distance_sensor(){
  pinMode(PIN_TRIG, OUTPUT);
  pinMode(PIN_ECHO, INPUT);
}


void setup_motor_drivers(){
  pinMode(USER_EMERGENCY_STOP_PIN, INPUT);
  pinMode(USER_DIRECTION_CHANGE_PIN, INPUT);
  pinMode(RGHT_WHEEL_STEP_PIN, OUTPUT);
  pinMode(LEFT_WHEEL_STEP_PIN, OUTPUT);
  pinMode(DIRECTION_PIN, OUTPUT);
  digitalWrite(DIRECTION_PIN, HIGH);
}


void measure_distance(){
  // Measure
  digitalWrite(PIN_TRIG, HIGH);
  digitalWrite(PIN_TRIG, LOW);

  // Process
  int duration = pulseIn(PIN_ECHO, HIGH);
  float distance = duration / 58;
  collision_check = (distance <= distance_threshold);
}


void check_emergency(){
  emergency_stopped = (digitalRead(USER_EMERGENCY_STOP_PIN) ? true : false);
}


void determine_direction(){
  current_dir = (digitalRead(USER_DIRECTION_CHANGE_PIN) ? -1 : 1);
  digitalWrite(DIRECTION_PIN, (current_dir == 1 ? HIGH : LOW));
}


void step_wheels(){
  timer1_delay_running = (millis() - timer1_delay_start < timer1_total_interval);
  if(timer1_delay_running) return;

  measure_distance();
  check_emergency();
  if(emergency_stopped || collision_check) return;

  determine_direction();

  // Move motors
  digitalWrite(RGHT_WHEEL_STEP_PIN, HIGH);
  digitalWrite(LEFT_WHEEL_STEP_PIN, HIGH);
  digitalWrite(RGHT_WHEEL_STEP_PIN, LOW);
  digitalWrite(LEFT_WHEEL_STEP_PIN, LOW);

  // Update variables
  total_steps++;
  current_step_pos += current_dir;
  total_travelled = total_steps * steps_to_distance;
  current_pos_mtr = current_step_pos * steps_to_distance;
}


void update_oled_display(){
  timer2_delay_running = (millis() - timer2_delay_start < timer2_total_interval);
  if(timer2_delay_running) return;

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println(String("Travel: " + String(total_travelled, 3)));
  display.setCursor(0, 10);
  display.println(String("Position: " + String(current_pos_mtr, 3)));
  display.setCursor(0, 20);
  display.println(String("Direction: " + String(current_dir)));
  display.setCursor(0, 30);
  display.println(String("Emergency Stopped? " + String(emergency_stopped)));
  display.setCursor(0, 40);
  display.println(String("Collision Check? " + String(collision_check)));
  display.setCursor(0, 50);
  display.display();
}


void setup() {
  // Set up OLED
  setup_oled();

  // Set up components
  setup_distance_sensor();
  setup_motor_drivers();
}


void loop() {
  step_wheels();
  update_oled_display();
}