/*******************************************************************************
 * Mecanum Robot Control Sketch for Arduino Uno R4 WiFi
 *
 * This sketch provides full control for a 4-wheel Mecanum robot. It uses an
 * Adafruit Motor Shield v2 for motor control, an Adafruit BNO055 IMU for
 * orientation sensing, and reads four quadrature encoders using interrupts.
 *
 * HARDWARE:
 * - Arduino Uno R4 WiFi
 * - Adafruit Motor/Stepper/Servo Shield for Arduino v2
 * - 4x DC Motors with Quadrature Encoders
 * - Adafruit BNO055 9-DOF Absolute Orientation IMU
 *
 * LIBRARIES REQUIRED:
 * - Adafruit Motor Shield V2 Library:
 * https://github.com/adafruit/Adafruit_Motor_Shield_V2_Library
 * - Adafruit BNO055 Library:
 * https://github.com/adafruit/Adafruit_BNO055
 * - Adafruit Unified Sensor Library:
 * https://github.com/adafruit/Adafruit_Sensor
 *
 * WIRING:
 * 1. Motor Shield:
 * - Sits directly on top of the Arduino Uno.
 * - Connect motors to M1, M2, M3, M4 terminals.
 * - M1: Front Left (fl)
 * - M2: Front Right (fr)
 * - M3: Rear Left (rl)
 * - M4: Rear Right (rr)
 * - Provide power to the shield's power terminal.
 *
 * 2. BNO055 IMU:
 * - Connect to the I2C pins (SDA, SCL).
 * - VIN to 5V, GND to GND.
 *
 * 3. Encoders:
 * - Connect encoder power to 5V and GND.
 * - Connect encoder A/B signal outputs to the specified digital pins.
 * - Encoder Front Left (fl):  Channel A -> D2, Channel B -> D5
 * - Encoder Front Right (fr): Channel A -> D7, Channel B -> D8
 * - Encoder Rear Left (rl):   Channel A -> D10, Channel B -> D11
 * - Encoder Rear Right (rr):  Channel A -> D12, Channel B -> D13
 *
 ******************************************************************************/

// Required libraries
#include <Wire.h>
#include <Adafruit_MotorShield.h>
#include <Adafruit_BNO055.h>
#include <utility/imumaths.h>

// --- PIN DEFINITIONS ---
// Encoders are mapped to specific interrupt-capable pins
#define ENCODER_FL_A 2
#define ENCODER_FL_B 5
#define ENCODER_FR_A 7
#define ENCODER_FR_B 8
#define ENCODER_RL_A 10
#define ENCODER_RL_B 11
#define ENCODER_RR_A 12
#define ENCODER_RR_B 13

// --- OBJECT INITIALIZATION ---
// Create the motor shield object with the default I2C address
Adafruit_MotorShield motor_shield = Adafruit_MotorShield();

// Assign motors to shield ports
Adafruit_DCMotor *motor_fl = motor_shield.getMotor(1); // Front Left
Adafruit_DCMotor *motor_fr = motor_shield.getMotor(2); // Front Right
Adafruit_DCMotor *motor_rl = motor_shield.getMotor(3); // Rear Left
Adafruit_DCMotor *motor_rr = motor_shield.getMotor(4); // Rear Right

// Create the BNO055 IMU object
// Use -1 for the ID to let the library auto-detect it
Adafruit_BNO055 bno = Adafruit_BNO055(-1, 0x28, &Wire);


// --- ENCODER STATE VARIABLES ---
// Volatile variables are used because they are modified inside ISRs
volatile long encoder_count_fl = 0;
volatile long encoder_count_fr = 0;
volatile long encoder_count_rl = 0;
volatile long encoder_count_rr = 0;

// These store the previous state of the A/B signals to decode direction
volatile uint8_t last_state_fl;
volatile uint8_t last_state_fr;
volatile uint8_t last_state_rl;
volatile uint8_t last_state_rr;

// Quadrature Encoder state transition table.
// This table provides a fast lookup to determine the direction of rotation.
// Index is calculated as: (last_state << 2) | current_state
// Values: 0 = no change, 1 = forward, -1 = backward
const int8_t qem_lookup_table[16] = {0, -1, 1, 0, 1, 0, 0, -1, -1, 0, 0, 1, 0, 1, -1, 0};


// --- INTERRUPT SERVICE ROUTINES (ISRs) ---
// These functions are executed automatically when an encoder pin changes state.
// They must be as fast as possible. No delays or serial prints inside.
// A separate ISR is required for each pin attached via attachInterrupt.

void isr_fl() {
  uint8_t current_state = (digitalRead(ENCODER_FL_A) << 1) | digitalRead(ENCODER_FL_B);
  uint8_t index = (last_state_fl << 2) | current_state;
  encoder_count_fl += qem_lookup_table[index];
  last_state_fl = current_state;
}

void isr_fr() {
  uint8_t current_state = (digitalRead(ENCODER_FR_A) << 1) | digitalRead(ENCODER_FR_B);
  uint8_t index = (last_state_fr << 2) | current_state;
  encoder_count_fr += qem_lookup_table[index];
  last_state_fr = current_state;
}

void isr_rl() {
  uint8_t current_state = (digitalRead(ENCODER_RL_A) << 1) | digitalRead(ENCODER_RL_B);
  uint8_t index = (last_state_rl << 2) | current_state;
  encoder_count_rl += qem_lookup_table[index];
  last_state_rl = current_state;
}

void isr_rr() {
  uint8_t current_state = (digitalRead(ENCODER_RR_A) << 1) | digitalRead(ENCODER_RR_B);
  uint8_t index = (last_state_rr << 2) | current_state;
  encoder_count_rr += qem_lookup_table[index];
  last_state_rr = current_state;
}


// --- SETUP FUNCTION ---
void setup() {
  Serial.begin(115200);
  Serial.println("Mecanum Robot Initializing...");

  // Initialize Motor Shield
  if (!motor_shield.begin()) {
    Serial.println("Could not find Motor Shield. Check connections.");
    while (1);
  }
  Serial.println("Motor Shield Initialized.");

  // Initialize BNO055 IMU
  if (!bno.begin()) {
    Serial.println("Could not find BNO055. Check connections.");
    while (1);
  }
  bno.setExtCrystalUse(true);
  Serial.println("BNO055 Initialized.");

  // Setup Encoder Pins
  pinMode(ENCODER_FL_A, INPUT_PULLUP);
  pinMode(ENCODER_FL_B, INPUT_PULLUP);
  pinMode(ENCODER_FR_A, INPUT_PULLUP);
  pinMode(ENCODER_FR_B, INPUT_PULLUP);
  pinMode(ENCODER_RL_A, INPUT_PULLUP);
  pinMode(ENCODER_RL_B, INPUT_PULLUP);
  pinMode(ENCODER_RR_A, INPUT_PULLUP);
  pinMode(ENCODER_RR_B, INPUT_PULLUP);

  // Get the initial state of the encoders
  last_state_fl = (digitalRead(ENCODER_FL_A) << 1) | digitalRead(ENCODER_FL_B);
  last_state_fr = (digitalRead(ENCODER_FR_A) << 1) | digitalRead(ENCODER_FR_B);
  last_state_rl = (digitalRead(ENCODER_RL_A) << 1) | digitalRead(ENCODER_RL_B);
  last_state_rr = (digitalRead(ENCODER_RR_A) << 1) | digitalRead(ENCODER_RR_B);

  // Attach interrupts for each encoder pin.
  // The ISR will be called on any CHANGE in the pin's state.
  attachInterrupt(digitalPinToInterrupt(ENCODER_FL_A), isr_fl, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_FL_B), isr_fl, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_FR_A), isr_fr, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_FR_B), isr_fr, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_RL_A), isr_rl, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_RL_B), isr_rl, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_RR_A), isr_rr, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_RR_B), isr_rr, CHANGE);

  Serial.println("Encoders Initialized with Interrupts.");
  Serial.println("Initialization Complete. Ready to move.");
  delay(1000);
}


// --- MAIN LOOP ---
void loop() {
  // This is a demonstration sequence of the robot's movements.
  // You can replace this with your own logic.

  Serial.println("Moving forward...");
  move_forward(200);
  delay(1000);
  stop_motors();
  delay(500);

  Serial.println("Moving backward...");
  move_backward(200);
  delay(1000);
  stop_motors();
  delay(500);

  Serial.println("Strafing right...");
  strafe_right(200);
  delay(1000);
  stop_motors();
  delay(500);

  Serial.println("Strafing left...");
  strafe_left(200);
  delay(1000);
  stop_motors();
  delay(500);

  Serial.println("Turning 90 degrees clockwise...");
  turn_relative(90.0, 150);
  delay(500);

  Serial.println("Turning 180 degrees counter-clockwise...");
  turn_relative(-180.0, 150);
  delay(500);

  Serial.println("Moving diagonal forward-right...");
  move_diagonal_forward_right(200);
  delay(1000);
  stop_motors();
  delay(500);

  Serial.println("Moving diagonal backward-left...");
  move_diagonal_backward_left(200);
  delay(1000);
  stop_motors();
  delay(500);

  // Print sensor data periodically
  print_sensor_data();

  Serial.println("Sequence complete. Restarting in 5 seconds.");
  delay(5000);
}

// --- HELPER FUNCTIONS ---

// Sets the speed and direction for all four motors.
void set_motor_speeds(int speed_fl, int speed_fr, int speed_rl, int speed_rr) {
  // Front Left Motor
  motor_fl->setSpeed(abs(speed_fl));
  if (speed_fl > 0) motor_fl->run(FORWARD);
  else if (speed_fl < 0) motor_fl->run(BACKWARD);
  else motor_fl->run(RELEASE);

  // Front Right Motor
  motor_fr->setSpeed(abs(speed_fr));
  if (speed_fr > 0) motor_fr->run(FORWARD);
  else if (speed_fr < 0) motor_fr->run(BACKWARD);
  else motor_fr->run(RELEASE);

  // Rear Left Motor
  motor_rl->setSpeed(abs(speed_rl));
  if (speed_rl > 0) motor_rl->run(FORWARD);
  else if (speed_rl < 0) motor_rl->run(BACKWARD);
  else motor_rl->run(RELEASE);

  // Rear Right Motor
  motor_rr->setSpeed(abs(speed_rr));
  if (speed_rr > 0) motor_rr->run(FORWARD);
  else if (speed_rr < 0) motor_rr->run(BACKWARD);
  else motor_rr->run(RELEASE);
}


// --- MOVEMENT FUNCTIONS ---

void stop_motors() {
  set_motor_speeds(0, 0, 0, 0);
}

void move_forward(uint8_t speed) {
  set_motor_speeds(speed, speed, speed, speed);
}

void move_backward(uint8_t speed) {
  set_motor_speeds(-speed, -speed, -speed, -speed);
}

void strafe_right(uint8_t speed) {
  set_motor_speeds(speed, -speed, -speed, speed);
}

void strafe_left(uint8_t speed) {
  set_motor_speeds(-speed, speed, speed, -speed);
}

void rotate_clockwise(uint8_t speed) {
  set_motor_speeds(speed, -speed, speed, -speed);
}

void rotate_counter_clockwise(uint8_t speed) {
  set_motor_speeds(-speed, speed, -speed, speed);
}

void move_diagonal_forward_right(uint8_t speed) {
  set_motor_speeds(speed, 0, 0, speed);
}

void move_diagonal_forward_left(uint8_t speed) {
  set_motor_speeds(0, speed, speed, 0);
}

void move_diagonal_backward_left(uint8_t speed) {
  set_motor_speeds(-speed, 0, 0, -speed);
}

void move_diagonal_backward_right(uint8_t speed) {
  set_motor_speeds(0, -speed, -speed, 0);
}


// --- IMU-BASED MOVEMENT ---

// Gets the current heading from the BNO055 (0-360 degrees)
float get_heading() {
  sensors_event_t event;
  bno.getEvent(&event);
  // The x-value of the orientation vector corresponds to heading
  return event.orientation.x;
}

// Turns the robot by a specific number of degrees relative to its current heading.
// Positive degrees for clockwise, negative for counter-clockwise.
void turn_relative(float angle_degrees, uint8_t speed) {
  float start_heading = get_heading();
  float target_heading = start_heading + angle_degrees;

  // Normalize target heading to be within 0-360
  if (target_heading >= 360.0) target_heading -= 360.0;
  if (target_heading < 0.0) target_heading += 360.0;

  Serial.print("Start Heading: "); Serial.println(start_heading);
  Serial.print("Target Heading: "); Serial.println(target_heading);

  // Start rotation
  if (angle_degrees > 0) {
    rotate_clockwise(speed);
  } else {
    rotate_counter_clockwise(speed);
  }

  // Loop until the target is reached
  while (true) {
    float current_heading = get_heading();
    float angle_turned = current_heading - start_heading;

    // Handle the 0/360 degree wrap-around for calculating turned angle
    if (angle_turned > 180.0) {
      angle_turned -= 360.0;
    } else if (angle_turned < -180.0) {
      angle_turned += 360.0;
    }

    // Check if we've turned far enough
    if (abs(angle_turned) >= abs(angle_degrees)) {
      break; // Exit loop
    }
    
    // Small delay to prevent busy-waiting
    delay(10); 
  }

  stop_motors();
  Serial.print("Turn complete. Final heading: ");
  Serial.println(get_heading());
}


// --- SENSOR DATA & DEBUGGING ---

void print_sensor_data() {
  // Get a new sensor event
  sensors_event_t orientation_data, ang_velocity_data, linear_accel_data;
  bno.getEvent(&orientation_data, Adafruit_BNO055::VECTOR_EULER);
  bno.getEvent(&ang_velocity_data, Adafruit_BNO055::VECTOR_GYROSCOPE);
  bno.getEvent(&linear_accel_data, Adafruit_BNO055::VECTOR_LINEARACCEL);

  Serial.println("--- IMU DATA ---");
  Serial.print("Heading (X): "); Serial.print(orientation_data.orientation.x, 2); Serial.println(" deg");
  Serial.print("Roll (Y):    "); Serial.print(orientation_data.orientation.y, 2); Serial.println(" deg");
  Serial.print("Pitch (Z):   "); Serial.print(orientation_data.orientation.z, 2); Serial.println(" deg");

  Serial.println("\n--- ENCODER COUNTS ---");
  Serial.print("Front Left:  "); Serial.println(encoder_count_fl);
  Serial.print("Front Right: "); Serial.println(encoder_count_fr);
  Serial.print("Rear Left:   "); Serial.println(encoder_count_rl);
  Serial.print("Rear Right:  "); Serial.println(encoder_count_rr);
  Serial.println("--------------------");
}