#include <Wire.h>
#include <VL53L0X.h>
#include <FastLED.h>

VL53L0X sensorUpward;
VL53L0X sensorForward;

#define NUM_LEDS 2
CRGB leds[NUM_LEDS];

#define M2A 17
#define M2B 16
#define M1A 19
#define M1B 18
#define WSLED 4

#define BUTTON1_PIN 21  // I01
#define BUTTON2_PIN 22  // I03

uint32_t lastTime = 0;

void errorLed(int count = 1){
  for (int i = 0; i < count; i++){
    leds[0] = CRGB::Red;
    FastLED.show();
    delay(100);    
    leds[0] = CRGB::Black;
    FastLED.show();
    delay(100);
  }
  delay(1000);
}

void reverse(uint8_t mySpeed){
  ledcWrite(3, mySpeed); //Set M2B to mySpeed
  ledcWrite(2, 0); //Set M2A to LOW
  ledcWrite(1, mySpeed); //Set M1B to mySpeed
  ledcWrite(0, 0); //Set M1A to LOW
}

void forward(uint8_t mySpeed){
  ledcWrite(2, mySpeed); //Set M2A to mySpeed
  ledcWrite(3, 0); //Set M2B to LOW
  ledcWrite(0, mySpeed); //Set M1A to mySpeed
  ledcWrite(1, 0); //Set M1B to LOW
}

void cw(uint8_t mySpeed){
  ledcWrite(3, mySpeed); //Set M2B to mySpeed
  ledcWrite(2, 0); //Set M2A to LOW
  ledcWrite(1, 0); //Set M1B to LOW
  ledcWrite(0, mySpeed); //Set M1A to mySpeed
}

void ccw(uint8_t mySpeed){
  ledcWrite(2, mySpeed); //Set M2A to mySpeed
  ledcWrite(3, 0); //Set M2B to LOW
  ledcWrite(0, 0); //Set M1A to LOW
  ledcWrite(1, mySpeed); //Set M1BA to mySpeed
}

void rotateLeft(uint8_t mySpeed){
  ccw(mySpeed); // Rotate left by counter-clockwise rotation
}

void rotateRight(uint8_t mySpeed){
  cw(mySpeed); // Rotate right by clockwise rotation
}

void setup() {
  Serial.begin(115200);
  FastLED.addLeds<NEOPIXEL, WSLED>(leds, NUM_LEDS);  // GRB ordering is assumed
  FastLED.setBrightness(24); // Set LED brightness to 50%

  // This section creates the I2C (Wire) connection and queries for a VL53LX
  Wire.begin();
  sensorUpward.setTimeout(500);
  sensorForward.setTimeout(500);
  if (!sensorUpward.init() || !sensorForward.init()) {
    Serial.println("Failed to detect and initialize sensor!");
    while (1) { errorLed(); }
  }

  // Set the 4 motor control lines as outputs
  pinMode(M1A, OUTPUT);
  pinMode(M1B, OUTPUT);
  pinMode(M2A, OUTPUT);
  pinMode(M2B, OUTPUT);

  // Set button pins as inputs with pull-up resistors
  pinMode(BUTTON1_PIN, INPUT_PULLUP);
  pinMode(BUTTON2_PIN, INPUT_PULLUP);

  // These lines connect 4 PWM channels to the 4 motor control Pins
  ledcSetup(0, 30000, 8);
  ledcAttachPin(M1A, 0);
  ledcSetup(1, 30000, 8);
  ledcAttachPin(M1B, 1);
  ledcSetup(2, 30000, 8);
  ledcAttachPin(M2A, 2);
  ledcSetup(3, 30000, 8);
  ledcAttachPin(M2B, 3);

  forward(230);
}

void loop() {
  int distanceForward = sensorForward.readRangeSingleMillimeters();
  int distanceUpward = sensorUpward.readRangeSingleMillimeters();

  // Check if the robot is under cover
  if (distanceUpward < 200) {
    // Robot is under cover, shake in place
    showSadEmotion();
  } else {
    // Cover is removed, move randomly
    moveRandomly(distanceForward);
  }

  // Check button states
  if (digitalRead(BUTTON1_PIN) == LOW) {
    showHappyEmotion();
    while (digitalRead(BUTTON1_PIN) == LOW); // Debounce
  } else if (digitalRead(BUTTON2_PIN) == LOW) {
    changeLEDs(CRGB::Red);
    while (digitalRead(BUTTON2_PIN) == LOW); // Debounce
  }

  Serial.print("Forward Distance: ");
  Serial.print(distanceForward);
  Serial.print(" mm, Upward Distance: ");
  Serial.print(distanceUpward);
  Serial.println(" mm");
}

void showSadEmotion() {
  // Shake motors
  for (int i = 0; i < 10; i++) {
    rotateLeft(200);
    delay(100);
    rotateRight(200);
    delay(100);
  }

  // Set LEDs to blue to indicate sadness
  leds[0] = CRGB::Blue;
  leds[1] = CRGB::Blue;
  FastLED.show();
}

void showHappyEmotion() {
  // LED Animation
  animateLEDs();

  // Dance Sequence
  danceSequence();

  // Figure Eight Movement
  figureEight();

  // Tail Wag
  tailWag();

  // Stop the motors after the happy emotion display
  ledcWrite(0, 0);
  ledcWrite(1, 0);
  ledcWrite(2, 0);
  ledcWrite(3, 0);
}

void animateLEDs() {
  for (int i = 0; i < 5; i++) {
    leds[0] = CRGB::Yellow;
    leds[1] = CRGB::Yellow;
    FastLED.show();
    delay(200);
    leds[0] = CRGB::Green;
    leds[1] = CRGB::Green;
    FastLED.show();
    delay(200);
  }
}

void danceSequence() {
  for (int i = 0; i < 3; i++) {
    forward(200);
    delay(500);
    cw(200);
    delay(500);
    ccw(200);
    delay(500);
    reverse(200);
    delay(500);
  }
}

void figureEight() {
  for (int i = 0; i < 2; i++) {
    forward(200);
    delay(1000);
    cw(200);
    delay(1000);
    forward(200);
    delay(1000);
    ccw(200);
    delay(1000);
  }
}

void tailWag() {
  for (int i = 0; i < 5; i++) {
    cw(200);
    delay(200);
    ccw(200);
    delay(200);
  }
}

void changeLEDs(CRGB color) {
  leds[0] = color;
  leds[1] = color;
  FastLED.show();
}

void moveRandomly(int distanceForward) {
  // Generate a random movement direction
  int randomDirection = random(0, 4);
  int randomDuration = random(300, 700); // Random duration for the movement

  switch (randomDirection) {
    case 0:
      forward(200);
      break;
    case 1:
      reverse(200);
      break;
    case 2:
      cw(200);
      break;
    case 3:
      ccw(200);
      break;
  }
  delay(randomDuration);

  // Avoid obstacles
  if (distanceForward < 50) {
    reverse(255);
    delay(200);
    cw(255);
    delay(150);
  } else if (distanceForward < 200) {
    forward(230 - (30 - ((distanceForward / 200.0) * 30)));
  } else {
    forward(230);
  }
}