// robot_controller_v3.ino
// Autonomous roomba-style navigation. Pi handles vision only (DROPL / DROPR).

#include "FastIMU.h"
#include "Adafruit_TCS34725.h"
#include <FastLED.h>
#include <Wire.h>
#include <Servo.h>
#include <VL53L1X.h>

#define MAX_BLUE_TILE_VISITS   1
#define MAX_TOTAL_KITS         8
#define VICTIM_WALL_DISTANCE   15

#define TCS_VCC_PIN    2
#define IMU_ADDRESS    0x68

#define XSHUT_F    3
#define XSHUT_B    4
#define XSHUT_L    A0
#define XSHUT_R    A1

#define IN3_PIN    10
#define IN4_PIN    9
#define IN1_PIN    12
#define IN2_PIN    11
#define ENA_PIN    5

#define IN5_PIN    8
#define IN6_PIN    7
#define IN7_PIN    A2
#define IN8_PIN    A3
#define ENB_PIN    6

#define SERVO_PIN          13
#define SERVO_STOP_US      1500
#define SERVO_FORWARD_US   1700
#define SERVO_BACKWARD_US  1300
#define SERVO_STEP_TIME_MS 1500

#define NUM_LEDS           25

#define YAW_CORRECTION_MULT  20
#define TURN_SPEED           200
#define TURN_STOP_LEAD       8
#define DRIVE_SPEED          120
#define REVERSE_SPEED        100
#define REVERSE_TIME_MS      700
#define MIN_WALL_DISTANCE    10

calData calib = { 0 };
GyroData gyroData;
MPU6050 IMU;
float yaw = 0;
float yaw_bias = 0;
uint32_t old_time;
bool imu_initialized = false;
bool initialized = false;
const int SAMPLE_AMOUNT = 400;

Adafruit_TCS34725 tcs = Adafruit_TCS34725(TCS34725_INTEGRATIONTIME_50MS, TCS34725_GAIN_4X);
CRGB leds[NUM_LEDS];
Servo medkit_dropper;
VL53L1X sensor_front;
VL53L1X sensor_back;
VL53L1X sensor_left;
VL53L1X sensor_right;

float target_yaw = 0;

int  blueTileCounter   = 0;
int  totalKitsDeployed = 0;
bool ignoreBlueTile    = false;
bool driving_forward   = false;

enum TileColor { TILE_WHITE, TILE_BLACK, TILE_BLUE };

struct raspberry_command {
  String function_name;
  int    value;
};

raspberry_command parse_input(const String& data_to_parse) {
  raspberry_command result;
  int divider = data_to_parse.indexOf(":");
  result.function_name = data_to_parse.substring(0, divider);
  if (divider == (int)data_to_parse.length() - 1) {
    result.value = 0;
  } else {
    result.value = data_to_parse.substring(divider + 1).toInt();
  }
  return result;
}

float read_laser_cm(VL53L1X &sensor) {
  int mm = sensor.readRangeContinuousMillimeters();
  if (sensor.timeoutOccurred() || mm <= 0 || mm > 4000) {
    return 999.0;
  }
  return mm / 10.0;
}

TileColor get_current_tile_color() {
  uint16_t r, g, b, c;
  tcs.getRawData(&r, &g, &b, &c);
  if (c == 0) {
    return TILE_WHITE;
  }
  if (c < 150) {
    return TILE_BLACK;
  }
  float relB = (float)b / c;
  float relR = (float)r / c;
  float relG = (float)g / c;
  if (relB > relR && relB > relG && relB > 0.45) {
    return TILE_BLUE;
  }
  return TILE_WHITE;
}

void stop_motors() {
  digitalWrite(IN3_PIN, LOW); digitalWrite(IN4_PIN, LOW);
  digitalWrite(IN1_PIN, LOW); digitalWrite(IN2_PIN, LOW);
  digitalWrite(IN5_PIN, LOW); digitalWrite(IN6_PIN, LOW);
  digitalWrite(IN7_PIN, LOW); digitalWrite(IN8_PIN, LOW);
  analogWrite(ENA_PIN, 0);   analogWrite(ENB_PIN, 0);
}

void set_direction_forward() {
  digitalWrite(IN3_PIN, LOW);  digitalWrite(IN4_PIN, HIGH);
  digitalWrite(IN1_PIN, LOW);  digitalWrite(IN2_PIN, HIGH);
  digitalWrite(IN5_PIN, LOW);  digitalWrite(IN6_PIN, HIGH);
  digitalWrite(IN7_PIN, HIGH); digitalWrite(IN8_PIN, LOW);
}

void set_direction_reverse() {
  digitalWrite(IN3_PIN, HIGH); digitalWrite(IN4_PIN, LOW);
  digitalWrite(IN1_PIN, HIGH); digitalWrite(IN2_PIN, LOW);
  digitalWrite(IN5_PIN, HIGH); digitalWrite(IN6_PIN, LOW);
  digitalWrite(IN7_PIN, LOW);  digitalWrite(IN8_PIN, HIGH);
}

void kickstart() {
  analogWrite(ENA_PIN, 255);
  analogWrite(ENB_PIN, 255);
  delay(50);
}

void update_yaw() {
  if (!imu_initialized) {
    return;
  }
  unsigned long now_time = micros();
  float passed_time = (now_time - old_time) / 1000000.0;
  old_time = now_time;
  IMU.update();
  IMU.getGyro(&gyroData);
  float gz = gyroData.gyroZ - yaw_bias;
  if (abs(gz) < 0.2) {
    gz = 0;
  }
  yaw += gz * passed_time;
}

float get_yaw() {
  update_yaw();
  return yaw;
}

void turn_left(int angle) {
  digitalWrite(IN3_PIN, LOW);  digitalWrite(IN4_PIN, HIGH);
  digitalWrite(IN1_PIN, LOW);  digitalWrite(IN2_PIN, HIGH);
  digitalWrite(IN5_PIN, HIGH); digitalWrite(IN6_PIN, LOW);
  digitalWrite(IN7_PIN, LOW);  digitalWrite(IN8_PIN, HIGH);
  float turn_target = get_yaw() + angle;
  kickstart();
  while (abs(turn_target - get_yaw()) > TURN_STOP_LEAD) {
    analogWrite(ENA_PIN, TURN_SPEED);
    analogWrite(ENB_PIN, TURN_SPEED);
  }
  stop_motors();
  digitalWrite(IN3_PIN, HIGH); digitalWrite(IN4_PIN, LOW);
  digitalWrite(IN1_PIN, HIGH); digitalWrite(IN2_PIN, LOW);
  digitalWrite(IN5_PIN, LOW);  digitalWrite(IN6_PIN, HIGH);
  digitalWrite(IN7_PIN, HIGH); digitalWrite(IN8_PIN, LOW);
  analogWrite(ENA_PIN, TURN_SPEED); analogWrite(ENB_PIN, TURN_SPEED);
  delay(25);
  stop_motors();
}

void turn_right(int angle) {
  digitalWrite(IN3_PIN, HIGH); digitalWrite(IN4_PIN, LOW);
  digitalWrite(IN1_PIN, HIGH); digitalWrite(IN2_PIN, LOW);
  digitalWrite(IN5_PIN, LOW);  digitalWrite(IN6_PIN, HIGH);
  digitalWrite(IN7_PIN, HIGH); digitalWrite(IN8_PIN, LOW);
  float turn_target = get_yaw() - angle;
  kickstart();
  while (abs(turn_target - get_yaw()) > TURN_STOP_LEAD) {
    analogWrite(ENA_PIN, TURN_SPEED);
    analogWrite(ENB_PIN, TURN_SPEED);
  }
  stop_motors();
  digitalWrite(IN3_PIN, LOW);  digitalWrite(IN4_PIN, HIGH);
  digitalWrite(IN1_PIN, LOW);  digitalWrite(IN2_PIN, HIGH);
  digitalWrite(IN5_PIN, HIGH); digitalWrite(IN6_PIN, LOW);
  digitalWrite(IN7_PIN, LOW);  digitalWrite(IN8_PIN, HIGH);
  analogWrite(ENA_PIN, TURN_SPEED); analogWrite(ENB_PIN, TURN_SPEED);
  delay(25);
  stop_motors();
}

void move_reverse() {
  set_direction_reverse();
  analogWrite(ENA_PIN, REVERSE_SPEED);
  analogWrite(ENB_PIN, REVERSE_SPEED);
  delay(REVERSE_TIME_MS);
  stop_motors();
}

void do_drop(int amount) {
  int actualDrop = amount;
  if (totalKitsDeployed + actualDrop > MAX_TOTAL_KITS) {
    actualDrop = MAX_TOTAL_KITS - totalKitsDeployed;
  }
  if (actualDrop < 0) {
    actualDrop = 0;
  }
  for (int i = 0; i < actualDrop; i++) {
    medkit_dropper.writeMicroseconds(SERVO_FORWARD_US);
    delay(SERVO_STEP_TIME_MS);
    medkit_dropper.writeMicroseconds(SERVO_STOP_US);
    delay(200);
    medkit_dropper.writeMicroseconds(SERVO_BACKWARD_US);
    delay(SERVO_STEP_TIME_MS);
    medkit_dropper.writeMicroseconds(SERVO_STOP_US);
    delay(200);
  }
  totalKitsDeployed += actualDrop;
  medkit_dropper.writeMicroseconds(SERVO_STOP_US);
}

void drop_medkit_left(int amount) {
  if (!initialized) {
    Serial.println("NOT_READY");
    return;
  }
  stop_motors();
  driving_forward = false;
  float dist = read_laser_cm(sensor_left);
  if (dist < VICTIM_WALL_DISTANCE) {
    do_drop(amount);
  }
  Serial.print("DONE:");
  Serial.println(totalKitsDeployed);
}

void drop_medkit_right(int amount) {
  if (!initialized) {
    Serial.println("NOT_READY");
    return;
  }
  stop_motors();
  driving_forward = false;
  float dist = read_laser_cm(sensor_right);
  if (dist < VICTIM_WALL_DISTANCE) {
    do_drop(amount);
  }
  Serial.print("DONE:");
  Serial.println(totalKitsDeployed);
}

void init_laser(VL53L1X &sensor, int xshut_pin, int address, const char* name) {
  pinMode(xshut_pin, INPUT);
  delay(10);
  sensor.setTimeout(500);
  if (!sensor.init()) {
    Serial.print("CHYBA: Laser neodpoveda - ");
    Serial.println(name);
    digitalWrite(xshut_pin, LOW);
  } else {
    sensor.setAddress(address);
    sensor.setDistanceMode(VL53L1X::Long);
    sensor.startContinuous(50);
    Serial.print("OK: Laser ");
    Serial.println(name);
  }
}

void setup() {
  Serial.begin(115200);
  Serial.setTimeout(10);

  pinMode(TCS_VCC_PIN, OUTPUT);
  digitalWrite(TCS_VCC_PIN, LOW);

  pinMode(XSHUT_F, OUTPUT); pinMode(XSHUT_B, OUTPUT);
  pinMode(XSHUT_L, OUTPUT); pinMode(XSHUT_R, OUTPUT);
  digitalWrite(XSHUT_F, LOW); digitalWrite(XSHUT_B, LOW);
  digitalWrite(XSHUT_L, LOW); digitalWrite(XSHUT_R, LOW);
  delay(10);

  FastLED.addLeds<WS2812, 1>(leds, NUM_LEDS);
  FastLED.setBrightness(32);

  Wire.begin();
  Wire.setClock(50000);
  Wire.setWireTimeout(25000, true);

  medkit_dropper.attach(SERVO_PIN);
  medkit_dropper.writeMicroseconds(SERVO_STOP_US);

  pinMode(IN1_PIN, OUTPUT); pinMode(IN2_PIN, OUTPUT);
  pinMode(IN3_PIN, OUTPUT); pinMode(IN4_PIN, OUTPUT);
  pinMode(ENA_PIN, OUTPUT);
  pinMode(IN5_PIN, OUTPUT); pinMode(IN6_PIN, OUTPUT);
  pinMode(IN7_PIN, OUTPUT); pinMode(IN8_PIN, OUTPUT);
  pinMode(ENB_PIN, OUTPUT);

  delay(500);

  init_laser(sensor_front, XSHUT_F, 0x30, "PREDNY");
  init_laser(sensor_back,  XSHUT_B, 0x31, "ZADNY");
  init_laser(sensor_left,  XSHUT_L, 0x32, "LAVY");
  init_laser(sensor_right, XSHUT_R, 0x33, "PRAVY");

  digitalWrite(TCS_VCC_PIN, HIGH);
  delay(100);
  if (!tcs.begin()) {
    Serial.println("CHYBA: TCS34725 nabehol nespravne.");
  } else {
    Serial.println("OK: Senzor farby pripraveny.");
  }

  yaw = 0;
  int err = -1;
  int pokus = 0;
  while (err != 0 && pokus < 20) {
    err = IMU.init(calib, IMU_ADDRESS);
    delay(200);
    pokus++;
  }
  if (err != 0) {
    Serial.print("Error initializing IMU: ");
    Serial.println(err);
  } else {
    IMU.setGyroRange(500);
    double sum = 0;
    for (int i = 0; i < SAMPLE_AMOUNT; i++) {
      IMU.update();
      IMU.getGyro(&gyroData);
      sum += (double)gyroData.gyroZ;
      delay(5);
    }
    yaw_bias = (float)(sum / SAMPLE_AMOUNT);
  }

  imu_initialized = true;
  old_time = micros();
  initialized = true;
}

void loop() {
  if (Serial.available()) {
    String data = Serial.readStringUntil('\n');
    raspberry_command cmd = parse_input(data);
    if (cmd.function_name.equalsIgnoreCase("DROPL")) {
      drop_medkit_left(cmd.value);
    } else if (cmd.function_name.equalsIgnoreCase("DROPR")) {
      drop_medkit_right(cmd.value);
    }
  }

  if (!initialized) {
    return;
  }

  update_yaw();

  for (int i = 0; i < NUM_LEDS; i++) {
    leds[i] = CRGB::White;
  }
  FastLED.show();

  TileColor color = get_current_tile_color();

  if (color == TILE_BLACK) {
    stop_motors();
    driving_forward = false;
    move_reverse();
    float l = read_laser_cm(sensor_left);
    float r = read_laser_cm(sensor_right);
    if (l > r) {
      turn_left(90);
    } else {
      turn_right(90);
    }
    return;
  }

  if (color == TILE_BLUE && !ignoreBlueTile) {
    stop_motors();
    driving_forward = false;
    unsigned long startWait = millis();
    while (millis() - startWait < 5000) {
      if (Serial.available()) {
        String data = Serial.readStringUntil('\n');
        raspberry_command cmd = parse_input(data);
        if (cmd.function_name.equalsIgnoreCase("DROPL")) {
          drop_medkit_left(cmd.value);
        } else if (cmd.function_name.equalsIgnoreCase("DROPR")) {
          drop_medkit_right(cmd.value);
        }
      }
      delay(10);
    }
    blueTileCounter++;
    if (blueTileCounter >= MAX_BLUE_TILE_VISITS) {
      ignoreBlueTile = true;
    }
    return;
  }

  float front = read_laser_cm(sensor_front);
  if (front < MIN_WALL_DISTANCE) {
    stop_motors();
    driving_forward = false;
    float l = read_laser_cm(sensor_left);
    float r = read_laser_cm(sensor_right);
    if (l > r && l > MIN_WALL_DISTANCE) {
      turn_left(90);
    } else if (r >= l && r > MIN_WALL_DISTANCE) {
      turn_right(90);
    } else {
      move_reverse();
    }
    return;
  }

  if (!driving_forward) {
    target_yaw = get_yaw();
    set_direction_forward();
    kickstart();
    driving_forward = true;
  }

  float yaw_error = target_yaw - yaw;
  float correction = yaw_error * YAW_CORRECTION_MULT;
  int left_pwm = constrain(DRIVE_SPEED - (int)correction, 0, 255);
  int right_pwm = constrain(DRIVE_SPEED + (int)correction, 0, 255);
  analogWrite(ENA_PIN, right_pwm);
  analogWrite(ENB_PIN, left_pwm);
}
