// Libraries

#include <SparkFun_TB6612.h>
#include <Wire.h>
#include <Adafruit_Sensor.h>
#include <Arduino_Helpers.h>
#include <AH/Math/Quaternion.hpp>
#include <Adafruit_MPU6050.h>
#include "Adafruit_AHRS_Mahony.h"
#include <ArduinoBLE.h>

//MOTOR 
// Pins for all inputs, keep in mind the PWM defines must be on PWM pins
#define STBY D7
#define AIN1 D8
#define AIN2 D9
#define PWMA D10

// the following constant is used to allow you to make your motor configuration 
// line up with function names like forward.  Value can be 1 or -1 
const int offsetA = 1;

// Initializing the linear actuator  
Motor motor1 = Motor(AIN1, AIN2, PWMA, offsetA, STBY);

//STATE MACHINE
unsigned long sTime;
unsigned long cTime;

#define SM_START 0
#define SM_IMU 1
#define SM_FILL 2
#define SM_DRAIN 3
#define SM_MAINTAIN 4

uint8_t cstate;
uint8_t ostate;

// IMU
const int MPU1_addr = 0x68;  // I2C address of the first MPU6050 --> A4 + A5
const int MPU2_addr = 0x69;  // I2C address of the second MPU6050 --> A4 + A5

Adafruit_MPU6050 mpu1;
Adafruit_MPU6050 mpu2;

Quaternion q1, q2;

Adafruit_Mahony filter1;
Adafruit_Mahony filter2;

//Wirst position
bool wristExtended = false;  // Flag indicating whether the wrist is flexed or extended

const int ThresholdH = 35;  // Threshold for wrist flexion/extension detection
const int ThresholdL = -35;  // Threshold for wrist flexion/extension detection

//Moving Average Filter
const int numReadings = 10;             // Number of readings for moving average
float DifferenceReadings[numReadings];  // Array to store angle differences
int currentIndex = 0;                   // Index to keep track of the current reading position
float filteredDifferenceRoll;
float filteredDifferenceYaw;
float filteredDifferencePitch;

//FILL and DRAIN variables
int flag;

//BUTTON
#define BUTTON_PIN 3
int status;
int button_state_old = 0;
int flag_button = 0;

//----------------------------------------------------------------------
// Voltage divider definitions. 
// Here I've reported my precise values for R1 (2.66K) and R2 (1K)
// remember: you cannot apply more than 3.3V on analog input. 
// Having a voltage divider with those values (2.66K and 1K)
// consent to apply on the voltage divider input a maximum of:
// 3.3/VDK = 3.3/0.275 = 12V
#define R1  2660
#define R2  1000
#define VDK  0.2732240437//(R2/(R1+R2))

// since I'm using a Li-Ion battery 11.1V and you cannot apply
// less than 5V on Vin pin, I define my "battery full" value as 11.1V
// and "battery low" value as 5V
#define BAT_FUL 11.1
#define BAT_LOW 5

#define BAT_DIS 20 //percentage
#define BAT_CHARG 12 //voltage while connected to the line

// analog pin connected to the voltage divider for battery 
// voltage reading
#define AN_BAT  A3

//LED defines
#define RED_PIN  D6
#define GREEN_PIN  D5
#define BLUE_PIN  D4

static long preMillis = 0;
uint16_t b_averages=100; // average for battery
  
// battery averaged value/actual value
static float anBa=0;
float anBa1; 

static uint16_t i_b=0; // battery readings counter
static int flag_start = 1;

//--------------------------------------------------------------------------

void init_sm() {
  cstate = SM_START;
}

void update_sm() {

  switch (cstate) {
    case SM_START:
      Serial.println("\nSM_CALIBRATION");
      ostate = cstate;
      cstate = SM_IMU;
      break;

    case SM_IMU: 
      status = digitalRead(BUTTON_PIN);
      if (button_state_old == 0 && status){
          button_state_old = 1;
          flag_button = 1;
        }
      if(button_state_old == 1 && !status){        
          button_state_old = 0;
          flag_button = 1;}

      if (((filteredDifferenceRoll < ThresholdL) || flag_button) && flag) {
        // Svuoto per un periodo
        flag_button = 0;
        Serial.println("\nSM_DRAIN");
        sTime = millis();
        ostate = cstate;
        cstate = SM_DRAIN;
      }

      if (((filteredDifferenceRoll > ThresholdH) || flag_button) && !flag) {
        flag_button = 0;        
        Serial.println("\nSM_FILL");
        sTime = millis();
        ostate = cstate;
        cstate = SM_FILL;
      }
      break;

    case SM_FILL:  

      flag=1;
      //Use of the drive function which takes as arguements the speed
      //and optional duration.  A negative speed will cause it to go
      //backwards.  Speed can be from -255 to 255.  Also use of the 
      //brake function which takes no arguements.
      motor1.drive(-255,3000);
      motor1.brake();
      Serial.println("\nSM_IMU");
      sTime = millis();
      ostate = cstate;
      cstate = SM_IMU;
      break;

    case SM_DRAIN:
    
      flag=0;
      motor1.drive(+255,3000);
      motor1.brake();
      Serial.println("\nSM_IMU");
      sTime = millis();
      ostate = cstate;
      cstate = SM_IMU;
      break;

  }
}


void setup() {
  //Initialize SM
  init_sm();

  //Touch Button
  pinMode(BUTTON_PIN, INPUT);

  //Battery LED
  pinMode(RED_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(BLUE_PIN, OUTPUT);

  // IMU
  Wire.begin();              // Initialize the I2C bus
  initializeMPU(MPU1_addr);  // Initialize the first MPU6050
  initializeMPU(MPU2_addr);  // Initialize the second MPU6050

  // Initialize the angle difference readings array
  for (int i = 0; i < numReadings; i++) {
    DifferenceReadings[i] = 0.0;
  }

  motor1.drive(255,3000);
  motor1.brake();

  //Serial
  Serial.begin(9600);  // opens serial port, sets data rate to 9600 bps
  
  setImu();
  updateValues();
}

void loop() {
  long curMillis = millis();
  if (preMillis>curMillis) preMillis=0; // millis() rollover?
  if (curMillis - preMillis >= 5) // check values every 5mS
  {
    preMillis = curMillis;

    //IMU----------------------------------------------------------------------
    /* Get a new sensor event */ 
    sensors_event_t a, g, temp;
    mpu1.getEvent(&a, &g, &temp);
    filter1.updateIMU(g.gyro.x, g.gyro.y, g.gyro.z, a.acceleration.x, a.acceleration.y, a.acceleration.z);
    mpu2.getEvent(&a, &g, &temp);
    filter2.updateIMU(g.gyro.x, g.gyro.y, g.gyro.z, a.acceleration.x, a.acceleration.y, a.acceleration.z);

    // Read IMU

    float w, x, y, z;

    filter1.getQuaternion(&w, &x, &y, &z);
    q1 = Quaternion(w, x, y, z);
    filter2.getQuaternion(&w, &x, &y, &z);
    q2 = Quaternion(w, x, y, z);

    Quaternion qprod = Quaternion().hamiltonianProduct(q1, q2.conjugated());

    EulerAngles angles = EulerAngles().quat2eul(qprod);
    const float piGreco = 3.14159;

    float yaw = (angles.yaw / (piGreco)) * 180;
    float roll = (angles.roll / (piGreco)) * 180;
    float pitch = (angles.pitch / (piGreco)) * 180;
    
    //String s = String(yaw) + "," + String(pitch) + "," + String(roll);
    //Serial.println(s);

    // Apply moving average filter to the acceleration difference
    //filteredDifferenceYaw = calculateMovingAverage(yaw);
    filteredDifferenceRoll = calculateMovingAverage(roll);
    //filteredDifferencePitch = calculateMovingAverage(pitch);
    //String s = String(filteredDifferenceRoll) + "," + String(roll);
    //String s = String(filteredDifferenceYaw) + "," + String(filteredDifferenceRoll) + "," + String(filteredDifferencePitch);
    //Serial.println(s);
    //----------------------------------------------------------------------------------
    updateValues(); // call function for updating Battery value
    update_sm();
  }
}


void setImu() {
  Serial.println("Adafruit MPU6050 init MPU1!");

  // Try to initialize!
  if (!mpu1.begin(0x68)) {
    Serial.println("Failed to find MPU6050 chip");
    while (1) {
      delay(10);
    }
  }
  Serial.println("MPU6050 Found!");

  mpu1.setAccelerometerRange(MPU6050_RANGE_8_G);
  Serial.print("Accelerometer range set to: ");
  switch (mpu1.getAccelerometerRange()) {
    case MPU6050_RANGE_2_G:
      Serial.println("+-2G");
      break;
    case MPU6050_RANGE_4_G:
      Serial.println("+-4G");
      break;
    case MPU6050_RANGE_8_G:
      Serial.println("+-8G");
      break;
    case MPU6050_RANGE_16_G:
      Serial.println("+-16G");
      break;
  }
  mpu1.setGyroRange(MPU6050_RANGE_500_DEG);
  Serial.print("Gyro range set to: ");
  switch (mpu1.getGyroRange()) {
    case MPU6050_RANGE_250_DEG:
      Serial.println("+- 250 deg/s");
      break;
    case MPU6050_RANGE_500_DEG:
      Serial.println("+- 500 deg/s");
      break;
    case MPU6050_RANGE_1000_DEG:
      Serial.println("+- 1000 deg/s");
      break;
    case MPU6050_RANGE_2000_DEG:
      Serial.println("+- 2000 deg/s");
      break;
  }

  mpu1.setFilterBandwidth(MPU6050_BAND_21_HZ);
  Serial.print("Filter bandwidth set to: ");
  switch (mpu1.getFilterBandwidth()) {
    case MPU6050_BAND_260_HZ:
      Serial.println("260 Hz");
      break;
    case MPU6050_BAND_184_HZ:
      Serial.println("184 Hz");
      break;
    case MPU6050_BAND_94_HZ:
      Serial.println("94 Hz");
      break;
    case MPU6050_BAND_44_HZ:
      Serial.println("44 Hz");
      break;
    case MPU6050_BAND_21_HZ:
      Serial.println("21 Hz");
      break;
    case MPU6050_BAND_10_HZ:
      Serial.println("10 Hz");
      break;
    case MPU6050_BAND_5_HZ:
      Serial.println("5 Hz");
      break;
  }

  Serial.println("Adafruit MPU6050 init MPU2!");

  // Try to initialize!
  if (!mpu2.begin(0x69)) {
    Serial.println("Failed to find MPU6050 chip");
    while (1) {
      delay(10);
    }
  }
  Serial.println("MPU6050 Found!");

  mpu2.setAccelerometerRange(MPU6050_RANGE_8_G);
  Serial.print("Accelerometer range set to: ");
  switch (mpu2.getAccelerometerRange()) {
    case MPU6050_RANGE_2_G:
      Serial.println("+-2G");
      break;
    case MPU6050_RANGE_4_G:
      Serial.println("+-4G");
      break;
    case MPU6050_RANGE_8_G:
      Serial.println("+-8G");
      break;
    case MPU6050_RANGE_16_G:
      Serial.println("+-16G");
      break;
  }
  mpu2.setGyroRange(MPU6050_RANGE_500_DEG);
  Serial.print("Gyro range set to: ");
  switch (mpu2.getGyroRange()) {
    case MPU6050_RANGE_250_DEG:
      Serial.println("+- 250 deg/s");
      break;
    case MPU6050_RANGE_500_DEG:
      Serial.println("+- 500 deg/s");
      break;
    case MPU6050_RANGE_1000_DEG:
      Serial.println("+- 1000 deg/s");
      break;
    case MPU6050_RANGE_2000_DEG:
      Serial.println("+- 2000 deg/s");
      break;
  }

  mpu2.setFilterBandwidth(MPU6050_BAND_21_HZ);
  Serial.print("Filter bandwidth set to: ");
  switch (mpu2.getFilterBandwidth()) {
    case MPU6050_BAND_260_HZ:
      Serial.println("260 Hz");
      break;
    case MPU6050_BAND_184_HZ:
      Serial.println("184 Hz");
      break;
    case MPU6050_BAND_94_HZ:
      Serial.println("94 Hz");
      break;
    case MPU6050_BAND_44_HZ:
      Serial.println("44 Hz");
      break;
    case MPU6050_BAND_21_HZ:
      Serial.println("21 Hz");
      break;
    case MPU6050_BAND_10_HZ:
      Serial.println("10 Hz");
      break;
    case MPU6050_BAND_5_HZ:
      Serial.println("5 Hz");
      break;
  }
  Serial.println("MPU OK!");
  delay(100);

  filter1.begin(5);
  filter2.begin(5);
}

void initializeMPU(int addr) {
  // Function to initialize the MPU6050 for accelerometer readings
  Wire.beginTransmission(addr);  // Start communication with the MPU6050
  Wire.write(0x6B);              // PWR_MGMT_1 register
  Wire.write(0);                 // Set to zero (wakes up the MPU-6050)
  Wire.endTransmission(true);

  Wire.beginTransmission(addr);
  Wire.write(0x1C);  // ACCEL_CONFIG register
  Wire.write(0x00);  // Set full-scale range to ±2g
  Wire.endTransmission(true);
}

void readAccelerometer(int addr, int16_t& ax, int16_t& ay, int16_t& az) {
  // Function to read accelerometer data from MPU6050
  Wire.beginTransmission(addr);
  Wire.write(0x3B);  // Start with register 0x3B (ACCEL_XOUT_H)
  Wire.endTransmission(false);
  Wire.requestFrom(addr, 6, true);  // Request 6 bytes from MPU6050

  ax = Wire.read() << 8 | Wire.read();  // Read the accelerometer data
  ay = Wire.read() << 8 | Wire.read();
  az = Wire.read() << 8 | Wire.read();
}

void readGyroscope(int addr, int16_t& gx, int16_t& gy, int16_t& gz) {
  // Function to read gyroscope data from MPU6050
  Wire.beginTransmission(addr);
  Wire.write(0x43);  // Start with register 0x43 (GYRO_XOUT_H)
  Wire.endTransmission(false);
  Wire.requestFrom(addr, 6, true);  // Request 6 bytes from MPU6050

  gx = Wire.read() << 8 | Wire.read();  // Read the gyroscope data
  gy = Wire.read() << 8 | Wire.read();
  gz = Wire.read() << 8 | Wire.read();
}


float calculateMovingAverage(float newValue) {
  // Function to calculate moving average of angle differences
  // Update the array with the new value
  DifferenceReadings[currentIndex] = newValue;

  // Increment the current index and wrap around if needed
  currentIndex++;
  if (currentIndex >= numReadings) {
    currentIndex = 0;
  }

  // Calculate the average
  float average = 0.0;
  for (int i = 0; i < numReadings; i++) {
    average += DifferenceReadings[i];
  }
  average /= numReadings;

  return average;
}


void updateValues() {
  // read battery value
  anBa1=analogRead(AN_BAT);
  anBa+=anBa1;
  i_b++;
  if (i_b==b_averages||flag_start)
    {
    flag_start = 0;
    Serial.println(anBa);
    anBa/=b_averages; // averaged analog value
    float voltage=anBa*(3.3/1023); // voltage on pin (if 3.3V => ADC gives 1023)
    voltage=voltage/VDK; // real voltage on the voltage divider input = battery voltage
    float Vin = voltage;
    
    Serial.print("Battery: ");
    Serial.print(voltage,2);
    Serial.print("V (");

    // calculate percentual battery
    // we must consider BAT_FUL=100% and BAT_LO=0%
    // report voltage to range having 0 as lower limit and (BAT_FUL-BAT_LO) as higher limit
    // for having percent value
    voltage-=BAT_LOW; 
    voltage/=(BAT_FUL-BAT_LOW); 
    anBa=voltage*100;
    // keep percent value in the range 0-100
    if (anBa<0) anBa=0;
    else if (anBa>100) anBa=100;
    
    Serial.print(anBa,1);
    Serial.println("%)");

    if(anBa < BAT_DIS){
      digitalWrite(RED_PIN, HIGH);   // Turn on the red LED
      digitalWrite(GREEN_PIN, LOW);  // Turn off the green LED
      digitalWrite(BLUE_PIN, LOW);  // Turn off the blue LED
    } else if (anBa >= BAT_DIS && Vin<BAT_CHAR) {
      digitalWrite(RED_PIN, LOW);    // Turn off the red LED
      digitalWrite(GREEN_PIN, HIGH); // Turn on the green LED
      digitalWrite(BLUE_PIN, LOW);  // Turn off the blue LED
    } else if(Vin>=BAT_CHARG){
      digitalWrite(RED_PIN, LOW);    // Turn off the red LED
      digitalWrite(GREEN_PIN, LOW); // Turn on the green LED
      digitalWrite(BLUE_PIN, HIGH);  // Turn off the blue LED
    }
        
    i_b=0;
    anBa=0;
    }
}


//MOTOR FUNCTIONS ------------------------------------------------------------------------------------------------------------------

// Drive in direction given by sign, at speed given by magnitude of the 
//parameter.
void drive(int speed);  

// drive(), but with a delay(duration)
void drive(int speed, int duration);  

//currently not implemented
//void stop();           // Stop motors, but allow them to coast to a halt.
//void coast();          // Stop motors, but allow them to coast to a halt.

//Stops motor by setting both input pins high
void brake(); 

//set the chip to standby mode.  The drive function takes it out of standby 
//(forward, back, left, and right all call drive)
void standby(); 

//Takes 2 motors and goes forward, if it does not go forward adjust offset 
//values until it does.  These will also take a negative number and go backwards
//There is also an optional speed input, if speed is not used, the function will
//use the DEFAULTSPEED constant.
void forward(Motor motor1, Motor motor2, int speed);
void forward(Motor motor1, Motor motor2);

//Similar to forward, will take 2 motors and go backwards.  This will take either
//a positive or negative number and will go backwards either way.  Once again the
//speed input is optional and will use DEFAULTSPEED if it is not defined.
void back(Motor motor1, Motor motor2, int speed);
void back(Motor motor1, Motor motor2);

//Left and right take 2 motors, and it is important the order they are sent.
//The left motor should be on the left side of the bot.  These functions
//also take a speed value
void left(Motor left, Motor right, int speed);
void right(Motor left, Motor right, int speed);

//This function takes 2 motors and and brakes them
void brake(Motor motor1, Motor motor2);