#include <Adafruit_MotorShield.h>

//General
uint8_t regions[2] = {BACKWARD, FORWARD};
uint8_t regions_opposite[2] = {FORWARD, BACKWARD}; 
int region = 1;
#define floating_avg_num 1
int button_val = 0;
int button_port = 0;

//Motors
Adafruit_MotorShield AFMS = Adafruit_MotorShield();
Adafruit_DCMotor *motors[] = {AFMS.getMotor(3), AFMS.getMotor(4)};
#define left 1 //index value in motor pointer array
#define right 0

//IR
#define white 1
#define black 0
#define num_IR 4
int IR_threshold[num_IR] = {200, 200, 1024, 1024};
float IR_avgs[num_IR] = {0}; //LRFB
int IR_colours[num_IR] = {0};
int IR_locns[num_IR] = {A0, A3, A1, A2};
float IR_vals[num_IR][floating_avg_num] = {0}; //check
int IR_float_counter[num_IR] = {0};

void setup() {
  // put your setup code here, to run once:
  Serial.begin(9600);
  if (!AFMS.begin()) {         // create with the default frequency 1.6KHz
  // if (!AFMS.begin(1000)) {  // OR with a different frequency, say 1KHz
    Serial.println("Could not find Motor Shield. Check wiring.");
    while (1);
  }
  for(int j=0; j<2; j++) {
    motors[j]->run(RELEASE);
  }
  for(int i=0; i< num_IR; i++) {
    pinMode(IR_locns[i], INPUT);
  }
  pinMode(button_port, INPUT);
}

void loop() {
  // put your main code here, to run repeatedly:
  for(int i=0; i < 2*floating_avg_num; i++) { //Get some values in before evaluation starts
    update_IR_vals();
  }
  for(int j = 0; j<2; j++) {
    motors[j]->run(regions[region]);
  }
  while (true) {
    if (button_val == 0) {
      update_IR_vals();
      int IR_check_val = IR_check(IR_colours);
      if (IR_check_val == 0) {
        move_both(motors, 255);
        Serial.println("FORWARD");
      }
      else if (IR_check_val == 1) {
        move_one_motor(motors, 200, left);
        Serial.println("TURN LEFT");
      }
      else if (IR_check_val == 2) {
        move_one_motor(motors, 200, right);
        Serial.println("TURN RIGHT");
      }
      else {
        Serial.println("Else block fml");
      }
      button_val = digitalRead(button_port);
    }
    else {
      Serial.println("Button pushed. Exiting");
      for(int j = 0; j < 2; j++) {
        motors[j]->run(RELEASE);
      }
      while (1);
    }  
  }
}
void move_one_motor(Adafruit_DCMotor *motors[], int velocity, int side) { //May need to account for direction later
  motors[(side+1)%2]->setSpeed(velocity);
  motors[side]->setSpeed(0);
}
void move_both(Adafruit_DCMotor *motors[], int velocity) {
  for(int j=0; j<2; j++) {
    motors[j]->setSpeed(velocity);
  }
}

float floating_average(int *index, float vals[], float curAvg, float newVal, int arr_len) {
  curAvg = ((curAvg * float(arr_len)) - vals[*index] + newVal)/float(arr_len);
  vals[*index] = newVal;
  *index = (*index + 1) % arr_len;
  return curAvg;
}

int IR_check(int LRFB[]) {
  if (LRFB[0] == black && LRFB[1] == black && LRFB[2] == white && LRFB[3] == white) {
    return 0;
  }
  else if (LRFB[0] == white && LRFB[1] == black) {
    return 1;
  }
  else if (LRFB[0] == black && LRFB[1] == white) {
    return 2;
  }
  else {
    return 3;
  }
}

void update_IR_vals() {
  for(int i=0; i<num_IR; i++) {
    int new_val = analogRead(IR_locns[i]);
    IR_avgs[i] = floating_average(&IR_float_counter[i], IR_vals[i], IR_avgs[i], new_val, floating_avg_num);
    if (IR_avgs[i] > IR_threshold[i]) {
      IR_colours[i] = black;
    }
    else {
      IR_colours[i] = white;
    }
  }
}