#include <ESP32Servo.h>
#include "FastIMU.h"
#include <Wire.h>

#define IMU_ADDRESS 0x68

BMI160 IMU;

calData calib = { 0 };
AccelData accelData;
GyroData gyroData;
MagData magData;

//     Arduino Nano ESP32
// #define CURRENT_1 A0
// #define CURRENT_2 A1
// #define SERVO_1 8
// #define SERVO_2 9

//          ESP32
#define CURRENT_1 34
#define CURRENT_2 35
#define SERVO_1 4
#define SERVO_2 16

uint8_t repetitions = 2;
uint8_t motors_speed = 100;
uint8_t new_repetitions = 2;
uint8_t new_motors_speed = 100;

uint8_t iterations;
uint8_t mode = 0;

uint8_t wave_num = 255;

Servo servo1;
Servo servo2;

void setup() {
  Wire.begin();
  Wire.setClock(400000);

  Serial0.begin(115200);
  Serial0.setTimeout(10);

  pinMode(CURRENT_1, INPUT);
  pinMode(CURRENT_2, INPUT);

  servo1.attach(SERVO_1);
  servo2.attach(SERVO_2);
  servo1.writeMicroseconds(1472);
  servo2.writeMicroseconds(1472);

  int err = IMU.init(calib, IMU_ADDRESS);
  if (err != 0) {
    Serial.print("Error initializing IMU: ");
    Serial.println(err);
    // while (true) {
    //   ;
    // }
  }

  Serial0.write(0xff);
  Serial0.write(0xff);
  Serial0.write(0xff);
}

void loop() {
  if (Serial0.available() > 0) {
    String data = Serial0.readStringUntil(';');
    if (data == "+rep_rel") {
      new_repetitions += 1;
      if (new_repetitions > 9) new_repetitions = 9;
      setStr("settings.t0", "txt", String(new_repetitions));
    } else if (data == "-rep_rel") {
      new_repetitions -= 1;
      if (new_repetitions < 1) new_repetitions = 1;
      setStr("settings.t0", "txt", String(new_repetitions));
    } else if (data == "+speed_rel") {
      new_motors_speed += 5;
      if (new_motors_speed > 100) new_motors_speed = 100;
      setStr("settings.t1", "txt", String(new_motors_speed));
    } else if (data == "-speed_rel") {
      new_motors_speed -= 5;
      if (new_motors_speed < 10) new_motors_speed = 10;
      setStr("settings.t1", "txt", String(new_motors_speed));
    }

    else if (data == "s-wave") {
      wave_num = 0;
    } else if (data == "w-next") {
      wave_num += 1;
      if (wave_num == 3) wave_num = 2;
      setStr("wave.t1", "txt", "Graph " + String(wave_num + 1));
    } else if (data == "w-last") {
      wave_num -= 1;
      if (wave_num == 255) wave_num = 0;
      setStr("wave.t1", "txt", "Graph " + String(wave_num + 1));
    }
    
    else if (data == "c-wave") wave_num = 255;

    else if (data == "p-settings") {
      setStr("settings.t0", "txt", String(repetitions));
      setStr("settings.t1", "txt", String(motors_speed));
    }

    else if (data == "settings-c") {
      new_motors_speed = motors_speed;
      new_repetitions = repetitions;
    } else if (data == "settings-s") {
      motors_speed = new_motors_speed;
      repetitions = new_repetitions;
      if (mode == 1) {
        setServoSpeed(motors_speed, iterations % 2);
      }
    }

    else if (data == "start") {
      start(1);
    } else if (data == "stop") {
      setServoSpeed(0, 0);
      mode = 0;
      return;
    }
  }

  handle_motors();
  handleCurrent();
  handleBMI();

  delay(10);
}

void start(uint8_t new_mode) {
  iterations = repetitions * 2;
  mode = new_mode;
}

void handle_motors() {
  static uint64_t tmr;

  if (millis() - tmr >= 2000) {
    if (mode == 1) {
      tmr = millis();
      if (iterations <= 0) {
        setServoSpeed(0, 0);
        mode = 0;
        return;
      }
      setServoSpeed(motors_speed, iterations % 2);
      iterations -= 1;
    }
  }
}

void handleCurrent() {
  static uint64_t tmr;

  if (millis() - tmr >= 15) {
    tmr = millis();

    if (wave_num == 0) {
      uint8_t I1 = map(analogRead(CURRENT_1), 0, 1023, 0, 255);
      uint8_t I2 = map(analogRead(CURRENT_2), 0, 1023, 0, 255);

      sendAdd(2, 0, I1);
      sendAdd(2, 1, I2);
    }
  }
}

void handleBMI() {
  static uint64_t tmr;

  if (millis() - tmr >= 30) {
    tmr = millis();

    IMU.update();
    IMU.getAccel(&accelData);
    IMU.getGyro(&gyroData);

    uint8_t mappedAccelX = map(accelData.accelX, 0, 15, 15, 240);
    uint8_t mappedAccelY = map(accelData.accelY, 0, 15, 15, 240);
    uint8_t mappedAccelZ = map(accelData.accelZ, 0, 15, 15, 240);
    if (wave_num == 1) {
      sendAdd(3, 0, mappedAccelX);
      sendAdd(3, 1, mappedAccelY);
      sendAdd(3, 2, mappedAccelZ);
    }

    uint8_t mappedGyroX = map(gyroData.gyroX, 0, 360, 15, 240);
    uint8_t mappedGyroY = map(gyroData.gyroY, 0, 360, 15, 240);
    uint8_t mappedGyroZ = map(gyroData.gyroZ, 0, 360, 15, 240);
    if (wave_num == 2) {
      sendAdd(4, 0, mappedGyroX);
      sendAdd(4, 1, mappedGyroY);
      sendAdd(4, 2, mappedGyroZ);
    }
  }
}

void setStr(String element, String param, String value) {
  Serial0.print(element);
  Serial0.print('.');
  Serial0.print(param);
  Serial0.print("=\"");
  Serial0.print(value);
  Serial0.print('"');
  Serial0.write(0xff);
  Serial0.write(0xff);
  Serial0.write(0xff);
}

void sendAdd(uint8_t id, uint8_t ch, uint8_t val) {
  Serial0.print("add ");
  Serial0.print(id);
  Serial0.print(',');
  Serial0.print(ch);
  Serial0.print(',');
  Serial0.print(val);
  Serial0.write(0xff);
  Serial0.write(0xff);
  Serial0.write(0xff);
}

void setServoSpeed(uint8_t speed, bool direction) {
  if (direction == 1) {
    int servo_micros = 1472 + 0.78 * speed;
    servo1.writeMicroseconds(servo_micros);
    servo2.writeMicroseconds(servo_micros);
  } else {
    int servo_micros = 1472 - 0.78 * speed;
    servo1.writeMicroseconds(servo_micros);
    servo2.writeMicroseconds(servo_micros);
  }
}