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

// ---------------------------------------------------------------------------
// Gyroskop
// ---------------------------------------------------------------------------

#define IMU_ADDRESS 0x68
calData calib = { 0 };
GyroData gyroData;
MPU6050 IMU;
float yaw;
float yaw_bias;
uint32_t old_time;
bool imu_initialized = false;
bool initialized = false;
const int SAMPLE_AMOUNT = 400;  // Priblizne 7.5ms na jeden

// ---------------------------------------------------------------------------
// Senzor farby
// ---------------------------------------------------------------------------

#define TCS_VCC_PIN 2
Adafruit_TCS34725 tcs = Adafruit_TCS34725(TCS34725_INTEGRATIONTIME_50MS, TCS34725_GAIN_4X);

// ---------------------------------------------------------------------------
// LED matica
// ---------------------------------------------------------------------------

#define NUM_LEDS 25
CRGB leds[NUM_LEDS];

// ---------------------------------------------------------------------------
// Servo
// ---------------------------------------------------------------------------

Servo medkit_dropper;
#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

// ---------------------------------------------------------------------------
// Laserove senzory vzdialenosti (4x VL53L1X)
// ---------------------------------------------------------------------------

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

VL53L1X sensor_front;
VL53L1X sensor_back;
VL53L1X sensor_left;
VL53L1X sensor_right;

// ---------------------------------------------------------------------------
// Motory
// ---------------------------------------------------------------------------

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

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

#define YAW_CORRECTION_MULT 20
#define TURN_SPEED 180

int current_speed = 0;
float target_yaw = 0;
float base_yaw = 0;

// ---------------------------------------------------------------------------
// Parsovanie prikazu
// ---------------------------------------------------------------------------

typedef void (*FunctionPointer)(int);

struct functions_table {
  const char* name;
  FunctionPointer function;
};

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 == data_to_parse.length() - 1) {
    result.value = 0;
  } else {
    result.value = data_to_parse.substring(divider + 1).toInt();
  }
  return result;
}

// ---------------------------------------------------------------------------
// Riadenie rychlosti
// ---------------------------------------------------------------------------

void apply_speed(int base_amount, float correction) {
  int left_pwm  = constrain((int)(abs(base_amount) - (int)correction), 0, 255);
  int right_pwm = constrain((int)(abs(base_amount) + (int)correction), 0, 255);

  if (base_amount > 0) {
    // Dozadu
    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);
  } else {
    // Dopredu
    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);
  }

  analogWrite(ENA_PIN, right_pwm);
  analogWrite(ENB_PIN, left_pwm);
}

void set_speed(int amount) {
  /*
  Nastavi rychlost robota. V rozmedzí od -100 do 100.
  Zaporne hodnoty = pohyb dozadu, kladne = pohyb dopredu, 0 = stop.
  */
  current_speed = map(amount, -100, 100, -255, 255);
  target_yaw = get_yaw();
  base_yaw = target_yaw;

  if (current_speed == 0) {
    analogWrite(ENA_PIN, 0);
    analogWrite(ENB_PIN, 0);
    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);
    return;
  }

  // Kratky rozbeh na prekonanie trenia v prevodovke
  int kickstart_amount = (current_speed < 0) ? -255 : 255;
  apply_speed(kickstart_amount, 0);
  delayMicroseconds(50);
  apply_speed(current_speed, 0);
}

// ---------------------------------------------------------------------------
// Otacanie
// ---------------------------------------------------------------------------

void turn_right(int angle) {
  // Prava strana dozadu, lava strana dopredu
  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;
  while (abs(turn_target - get_yaw()) > 2) {
    analogWrite(ENA_PIN, TURN_SPEED);
    analogWrite(ENB_PIN, TURN_SPEED);
  }
  set_speed(0);
  Serial.println("DONE");
}

void turn_left(int angle) {
  // Prava strana dopredu, lava strana dozadu
  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;
  while (abs(turn_target - get_yaw()) > 2) {
    analogWrite(ENA_PIN, TURN_SPEED);
    analogWrite(ENB_PIN, TURN_SPEED);
  }
  set_speed(0);
  Serial.println("DONE");
}

// ---------------------------------------------------------------------------
// Vypustenie medkitov
// ---------------------------------------------------------------------------

void drop_medkit(int amount) {
  set_speed(0);
  for (int i = 0; i < amount; 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);
  }
  medkit_dropper.writeMicroseconds(SERVO_STOP_US);
  Serial.println("DONE");
}

// ---------------------------------------------------------------------------
// Laserove senzory — pomocne funkcie
// ---------------------------------------------------------------------------

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

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);
    pinMode(xshut_pin, OUTPUT);
    digitalWrite(xshut_pin, LOW);
  } else {
    sensor.setAddress(address);
    sensor.setDistanceMode(VL53L1X::Long);
    sensor.startContinuous(50);
    Serial.print("OK: Laser "); Serial.println(name);
  }
}

// ---------------------------------------------------------------------------
// Senzory — hlavna funkcia
// ---------------------------------------------------------------------------

void get_sensors(int unused) {
  /*
  Vracia hodnoty zo senzorov v poradi:
  0. Predna vzdialenost (cm)
  1. Zadna vzdialenost (cm)
  2. Lava vzdialenost (cm)
  3. Prava vzdialenost (cm)
  4. R - farebny senzor
  5. G - farebny senzor
  6. B - farebny senzor
  7. Relativne otocenie (yaw)
  */
  float front_dist = read_laser_cm(sensor_front); delay(2);
  float back_dist  = read_laser_cm(sensor_back);  delay(2);
  float left_dist  = read_laser_cm(sensor_left);  delay(2);
  float right_dist = read_laser_cm(sensor_right); delay(2);

  int red, green, blue;
  uint16_t r, g, b, c;
  tcs.getRawData(&r, &g, &b, &c);

  if (c != 0) {
    float relR = (float)r / c;
    float relG = (float)g / c;
    float relB = (float)b / c;

    if (c < 150) {
      red = 0; green = 0; blue = 0;
    } else if (relB > relR && relB > relG && relB > 0.45) {
      red = 0; green = 0; blue = 255;
    } else {
      red = 255; green = 255; blue = 255;
    }
  } else {
    red = 255; green = 255; blue = 255;
  }

  Serial.print(front_dist); Serial.print(':');
  Serial.print(back_dist);  Serial.print(':');
  Serial.print(left_dist);  Serial.print(':');
  Serial.print(right_dist); Serial.print(':');
  Serial.print(red);        Serial.print(':');
  Serial.print(green);      Serial.print(':');
  Serial.print(blue);       Serial.print(':');
  Serial.println(yaw);
}

// ---------------------------------------------------------------------------
// IMU — aktualizacia a meranie
// ---------------------------------------------------------------------------

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 measure_YAW_bias(int samples, int delay_ms) {
  double sum = 0;
  for (int i = 0; i < samples; i++) {
    IMU.update();
    IMU.getGyro(&gyroData);
    sum += (double)gyroData.gyroZ;
    delay(delay_ms);
  }
  yaw_bias = (float)(sum / samples);
}

// ---------------------------------------------------------------------------
// Forward declaration
// ---------------------------------------------------------------------------

void start_system(int unused);

// ---------------------------------------------------------------------------
// Tabulka prikazov
// ---------------------------------------------------------------------------

functions_table functions[] = {
  { "SPEED",  set_speed },
  { "RIGHT",  turn_right },
  { "LEFT",   turn_left },
  { "DROP",   drop_medkit },
  { "SENSOR", get_sensors }
};

const int functions_count = sizeof(functions) / sizeof(functions[0]);

// ---------------------------------------------------------------------------
// Setup
// ---------------------------------------------------------------------------

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

  Serial.println("Inicializacia");

  // --- 1. TOTÁLNE VYPNUTIE SENZORA FARBY ---
  pinMode(TCS_VCC_PIN, OUTPUT);
  digitalWrite(TCS_VCC_PIN, LOW); // Odstrihneme TCS od prúdu

  // --- 2. XSHUT LASERY DOLE PRED Wire.begin() ---
  // VL53L1X senzory si pamätajú I2C adresu pokiaľ su napajane.
  // Musia byt v hardware resete (XSHUT=LOW) skor ako sa spusti I2C zbernica,
  // inak pri reinicializacii (napr. po odpojeni od PC) nie su na defaultnej adrese 0x29.
  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); // Daj sensorom cas na reset

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

  // --- 3. START I2C S NOVÝMI POISTKAMI ---
  Wire.begin();
  Wire.setClock(50000); // 50kHz pre oveľa stabilnejšiu komunikáciu
  Wire.setWireTimeout(25000, true); // Záchranná brzda: reset zbernice po 25ms chybe

  medkit_dropper.attach(SERVO_PIN);
  medkit_dropper.writeMicroseconds(SERVO_STOP_US);
  Serial.println("Motory");
  // Motory
  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);
  Serial.println("Lasery 1");
  // XSHUT piny su uz OUTPUT a LOW z kroku 2 — len pridame dlhsi reset
  delay(500);

  // --- 4. POSTUPNÉ ZAPNUTIE LASEROV ---
  Serial.println("Lasery 2");
  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");

  // --- 5. ZAPNUTIE SENZORA FARBY (TCS) ---
  Serial.println("Zapinam napajanie TCS...");
  digitalWrite(TCS_VCC_PIN, HIGH);
  delay(100);

  if (!tcs.begin()) {
    Serial.println("CHYBA: Senzor TCS34725 nabehol nespravne.");
  } else {
    Serial.println("OK: Senzor farby pripraveny.");
  }
  Serial.println("ZAcinam IMU");
  // --- 6. Inicializacia IMU ---
  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);
    measure_YAW_bias(SAMPLE_AMOUNT, 5);
  }

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

// ---------------------------------------------------------------------------
// Hlavna slucka
// ---------------------------------------------------------------------------

void loop() {
  if (Serial.available()) {
    String data_to_parse = Serial.readStringUntil('\n');
    raspberry_command parsed_data = parse_input(data_to_parse);

    for (int i = 0; i < functions_count; i++) {
      if (parsed_data.function_name.equalsIgnoreCase(functions[i].name)) {
        functions[i].function(parsed_data.value);
      }
    }
  }

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

  if (current_speed != 0 && imu_initialized) {
    float yaw_error = target_yaw - yaw;
    float correction = yaw_error * YAW_CORRECTION_MULT;
    apply_speed(current_speed, correction);
  }
}
