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

/* NAPAJANIE TCS SENZORA */
#define TCS_VCC_PIN 2 // Čip senzora farby

/*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;  // Nastaví sa na true po dokončení init_all
const int SAMPLE_AMOUNT = 400;  //Priblizne 7.5ms na jeden

/*Senzor farby*/
Adafruit_TCS34725 tcs = Adafruit_TCS34725(TCS34725_INTEGRATIONTIME_50MS, TCS34725_GAIN_4X);

/*Led matrix*/


/*Servo*/
Servo medkit_dropper;
#define SERVO_PIN 13
#define SERVO_STOP_US 1500
#define SERVO_FORWARD_US 1800
#define SERVO_BACKWARD_US 1200
#define SERVO_STEP_TIME_MS 1400

/*Laserove senzory vzdialenosti (4x VL53L1X)*/
#define XSHUT_F  5
#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  // Predny pravy motor
#define IN4_PIN 9   // Predny pravy motor
#define IN1_PIN 12  // Zadny pravy motor
#define IN2_PIN 11  // Zadny pravy motor
#define ENA_PIN 3   // PWM prava strana

// Lava strana
#define IN5_PIN 8   // Predny lavy motor
#define IN6_PIN 7   // Predny lavy motor
#define IN7_PIN A2  // Zadny lavy motor
#define IN8_PIN A3  // Zadny lavy motor
#define ENB_PIN 6   // PWM lava strana

#define YAW_CORRECTION_MULT 40
#define TURN_SPEED 200
#define TURN_TIMEOUT_MS 5000  // Bezpecnostny timeout pre otacanie
// Pocet stupnov pred cielom kde sa vypnu motory — momentum donesie robota na ciel
#define TURN_STOP_LEAD 8
#define TURN_SPEED_OFFSET 40
#define TURN_SPEED_SLOW 140       // Pomala rychlost pri jemnej korekcii (KALIBROVAT — nesmi zastavit motor!)
#define TURN_SLOW_THRESHOLD 25   // Stupne od cieloveho uhla, kde spomalime (KALIBROVAT)

/*Pohyb po dlazdiciach*/
#define TILE_SIZE 30
#define MOVE_MAX_SPEED       80
#define MOVE_MIN_SPEED_20     65  
#define MOVE_MIN_SPEED_10     60  
#define SLOWDOWN_KOEFICIENT   50  
#define MIN_WALL_DISTANCE 3.5

#define ENC_LEFT_PIN   0  
#define ENC_RIGHT_PIN  1   

#define ENCODER_PULSES_PER_REV  244.905

#define WHEEL_DIAMETER_MM       56.5f

const float DIST_PER_PULSE_MM = (3.14159265f * WHEEL_DIAMETER_MM) / ENCODER_PULSES_PER_REV;

volatile long enc_left_count  = 0;
volatile long enc_right_count = 0;

float target_yaw = 0;

typedef void (*FunctionPointer)(int);

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

// UPRAVENÁ FUNKCIA: Prečita vzdialenost z laseru a ošetrí chybné dáta
float read_laser_cm(VL53L1X &sensor) {
  // Serial.println("Meriam");
  int mm = sensor.readRangeContinuousMillimeters();
  
  // if (sensor.timeoutOccurred() || mm <= 0 || mm > 4000) {
  //   return 999.0;
  // }
  return mm / 10.0;
}

void stop_motors() {
  analogWrite(ENA_PIN, 0);
  analogWrite(ENB_PIN, 0);
}

void set_direction_reverse() {
  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, LOW); digitalWrite(IN8_PIN, HIGH);
}

void set_direction_forward() {
  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, HIGH); digitalWrite(IN8_PIN, LOW);
}

void turn_right(int angle) {
  stop_motors();
  delay(40);

  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, HIGH); digitalWrite(IN8_PIN, LOW);

  float turn_target = get_yaw() - angle;
  unsigned long turn_start = millis();

  // Faza 1: rychle otacanie, kym sme daleko od cieloveho uhla
  while (abs(turn_target - get_yaw()) > TURN_SLOW_THRESHOLD) {
    if (millis() - turn_start > TURN_TIMEOUT_MS) break;
    analogWrite(ENA_PIN, TURN_SPEED - TURN_SPEED_OFFSET);
    analogWrite(ENB_PIN, TURN_SPEED);
  }

  // Faza 2: pomale otacanie pre presnu korekciu
  while (abs(turn_target - get_yaw()) > TURN_STOP_LEAD) {
    if (millis() - turn_start > TURN_TIMEOUT_MS) break;
    analogWrite(ENA_PIN, TURN_SPEED_SLOW);
    analogWrite(ENB_PIN, TURN_SPEED_SLOW);
  }

  stop_motors();
  Serial.println("DONE");
}

void turn_left(int angle) {
  stop_motors();
  delay(40);

  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, LOW);  digitalWrite(IN8_PIN, HIGH);

  float turn_target = get_yaw() + angle;
  unsigned long turn_start = millis();

  // Faza 1: rychle otacanie, kym sme daleko od cieloveho uhla
  while (abs(turn_target - get_yaw()) > TURN_SLOW_THRESHOLD) {
    if (millis() - turn_start > TURN_TIMEOUT_MS) break;
    analogWrite(ENA_PIN, TURN_SPEED);
    analogWrite(ENB_PIN, TURN_SPEED - TURN_SPEED_OFFSET);
  }

  // Faza 2: pomale otacanie pre presnu korekciu
  while (abs(turn_target - get_yaw()) > TURN_STOP_LEAD) {
    if (millis() - turn_start > TURN_TIMEOUT_MS) break;
    analogWrite(ENA_PIN, TURN_SPEED_SLOW);
    analogWrite(ENB_PIN, TURN_SPEED_SLOW);
  }

  stop_motors();
  Serial.println("DONE");
}

void isr_enc_left()  { 
  enc_left_count++;  
}
void isr_enc_right() { 
  enc_right_count++; 
}

// Krátky impulz plného výkonu pre prekonanie statického trenia motorov
void kickstart() {
  analogWrite(ENA_PIN, 255);
  analogWrite(ENB_PIN, 255);
  delay(50);
}

void move_forward(int distance_cm) {
  target_yaw = get_yaw();
  set_direction_forward();
  Serial.println("resetrujem");
  noInterrupts();
  enc_left_count  = 0;
  enc_right_count = 0;
  interrupts();
  Serial.println("Doresetovane");

  long target_pulses = (long)((distance_cm * 10.0f) / DIST_PER_PULSE_MM);
  // Serial.println(target_pulses);
  int current_min_speed;
  if (distance_cm > 15) {
    current_min_speed = MOVE_MIN_SPEED_20;
  } else {
    current_min_speed = MOVE_MIN_SPEED_10;
  }

  // Kontrola pred pohybom — nespúšťame motor ak je pred nami stena
  float last_front_cm = read_laser_cm(sensor_front);
  if (last_front_cm < MIN_WALL_DISTANCE) {
    Serial.print("DONE:");
    Serial.println(0);
    return;
  }

  kickstart();

  long avg_pulses = 0;
  // Prvé I2C čítanie až po 50ms — last_front_cm je už inicializovaný platnou hodnotou
  unsigned long last_i2c_ms = millis();
  // Bezpecnostny timeout: max 5 sekund na akykolvek pohyb dopredu
  unsigned long move_start_ms = millis();
  while (avg_pulses < target_pulses) {
    noInterrupts();
    long lc = enc_left_count;
    long rc = enc_right_count;
    interrupts();
    avg_pulses = (lc + rc) / 2;

    if (millis() - move_start_ms > 5000) break;

    // I2C operacie (yaw + laser) max 1x za 50ms — VL53L1X je nastaveny na 50ms cyklus
    unsigned long now_ms = millis();
    if (now_ms - last_i2c_ms >= 50) {
      last_i2c_ms = now_ms;
      update_yaw();
      last_front_cm = read_laser_cm(sensor_front);
    }

    float d_remaining_mm = (target_pulses - avg_pulses) * DIST_PER_PULSE_MM;
    float t = constrain(d_remaining_mm / (SLOWDOWN_KOEFICIENT * 10.0f), 0.0f, 1.0f);
    int base_pwm = current_min_speed + (int)((MOVE_MAX_SPEED - current_min_speed) * t);
    float yaw_error = target_yaw - yaw;
    // Obmedzenie korekcie — zabranuje prudkemu zatacaniu pri poskodenom yaw
    float correction = constrain(yaw_error * YAW_CORRECTION_MULT, -10.0f, 10.0f);
    int left_pwm  = constrain(base_pwm - (int)correction, 0, 255);
    int right_pwm = constrain(base_pwm + (int)correction, 0, 255);
    analogWrite(ENA_PIN, right_pwm);
    analogWrite(ENB_PIN, left_pwm);

    if (last_front_cm < MIN_WALL_DISTANCE) break;
  }
  stop_motors();
  // Obnovenie merania predného senzora — po pohybe môže byť DRDY príznak zaseknutý
  sensor_front.startContinuous(50);

  noInterrupts();
  long final_avg = (enc_left_count + enc_right_count) / 2;
  interrupts();
  int actual_moved_cm = (int)(final_avg * DIST_PER_PULSE_MM / 10.0f);
  // Serial.print("Pulsy L:"); Serial.print(enc_left_count);
  // Serial.print(" R:"); Serial.println(enc_right_count);

  Serial.print("DONE:");
  Serial.println(actual_moved_cm);
}

void move_reverse(int distance_cm) {
  target_yaw = get_yaw();
  set_direction_reverse();

  noInterrupts();
  enc_left_count  = 0;
  enc_right_count = 0;
  interrupts();

  long target_pulses = (long)((distance_cm * 10.0f) / DIST_PER_PULSE_MM);
  int current_min_speed;
  if (distance_cm > 15) {
    current_min_speed = MOVE_MIN_SPEED_20;
  } else {
    current_min_speed = MOVE_MIN_SPEED_10;
  }

  // Kontrola pred pohybom — nespúšťame motor ak je za nami stena
  float last_back_cm = read_laser_cm(sensor_back);
  if (last_back_cm < MIN_WALL_DISTANCE) {
    Serial.println("DONE");
    return;
  }

  kickstart();

  long avg_pulses = 0;
  // Prvé I2C čítanie až po 50ms — last_back_cm je už inicializovaný platnou hodnotou
  unsigned long last_i2c_ms = millis();
  unsigned long move_start_ms = millis();
  while (avg_pulses < target_pulses) {
    noInterrupts();
    long lc = enc_left_count;
    long rc = enc_right_count;
    interrupts();
    avg_pulses = (lc + rc) / 2;

    if (millis() - move_start_ms > 5000) break;

    unsigned long now_ms = millis();
    if (now_ms - last_i2c_ms >= 50) {
      last_i2c_ms = now_ms;
      update_yaw();
      last_back_cm = read_laser_cm(sensor_back);
    }

    float d_remaining_mm = (target_pulses - avg_pulses) * DIST_PER_PULSE_MM;
    float t = constrain(d_remaining_mm / (SLOWDOWN_KOEFICIENT * 10.0f), 0.0f, 1.0f);
    int base_pwm = current_min_speed + (int)((MOVE_MAX_SPEED - current_min_speed) * t);
    float yaw_error = target_yaw - yaw;
    float correction = constrain(yaw_error * YAW_CORRECTION_MULT, -10.0f, 10.0f);
    int left_pwm  = constrain(base_pwm + (int)correction, 0, 255);
    int right_pwm = constrain(base_pwm - (int)correction, 0, 255);
    analogWrite(ENA_PIN, right_pwm);
    analogWrite(ENB_PIN, left_pwm);

    if (last_back_cm < MIN_WALL_DISTANCE) break;
  }
  stop_motors();
  // Obnovenie merania zadného senzora — po pohybe môže byť DRDY príznak zaseknutý
  sensor_back.startContinuous(50);

  Serial.println("DONE");
}

// Vypusti lievok dolava (rovnaky smer ako DROP)
void drop_left(int amount) {
  stop_motors();
  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");
}

// Vypusti lievok doprava (opacny smer ako DROP)
void drop_right(int amount) {
  stop_motors();
  for (int i = 0; i < amount; i++) {
    medkit_dropper.writeMicroseconds(SERVO_BACKWARD_US);
    delay(SERVO_STEP_TIME_MS);
    medkit_dropper.writeMicroseconds(SERVO_STOP_US);
    delay(200);
    medkit_dropper.writeMicroseconds(SERVO_FORWARD_US);
    delay(SERVO_STEP_TIME_MS);
    medkit_dropper.writeMicroseconds(SERVO_STOP_US);
    delay(200);
  }
  medkit_dropper.writeMicroseconds(SERVO_STOP_US);
  Serial.println("DONE");
}

// UPRAVENÁ FUNKCIA: Pridané delay(2) pre uvoľnenie I2C zbernice
void get_sensors(int unused) {
  float front_distance = read_laser_cm(sensor_front); delay(10);
  float back_distance = read_laser_cm(sensor_back);   delay(10);
  float left_distance = read_laser_cm(sensor_left);   delay(10);
  float right_distance = read_laser_cm(sensor_right); delay(10);
  
  int red, green, blue;
  uint16_t r, g, b, c;
  tcs.getRawData(&r, &g, &b, &c); delay(10);

  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_distance); Serial.print(':');
  Serial.print(back_distance); Serial.print(':');
  Serial.print(left_distance); Serial.print(':');
  Serial.print(right_distance); Serial.print(':');
  Serial.print(red); Serial.print(':');
  Serial.print(green); Serial.print(':');
  Serial.print(blue); Serial.print(':');
  Serial.println(yaw);
}

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);
}

functions_table functions[] = {
  {"RIGHT", turn_right},
  {"LEFT", turn_left},
  {"DROP_LEFT", drop_left},
  {"DROP_RIGHT", drop_right},
  {"SENSOR", get_sensors},
  {"FORWARD", move_forward},
  {"REVERSE", move_reverse},
  {"START",init_all}
};
const int functions_count = sizeof(functions) / sizeof(functions[0]);

void init_laser(VL53L1X &sensor, int xshut_pin, int address, const char* name) {
  Serial.println("ZApinam sensory");
  pinMode(xshut_pin,INPUT);
  delay(10);
  
  sensor.setTimeout(500);
  Serial.println("ZApinam sensory2");
  if (!sensor.init()) {
    Serial.print("CHYBA: Laser neodpoveda - ");
    Serial.println(name);
    digitalWrite(xshut_pin, LOW); // Vypnutie senzora, aby nebránil adrese ostatných
  } else {
    sensor.setAddress(address);
    sensor.setDistanceMode(VL53L1X::Long);
    sensor.startContinuous(50);
    Serial.print("OK: Laser "); Serial.println(name);
  }
  Serial.println("ZApinam sensory3");
}

// ---------------------------------------------------------------------------
// Inicializácia všetkých periférií
// ---------------------------------------------------------------------------
void init_all(int unused) {
  if (initialized){
    Serial.println("DONE");
  } else {
    Serial.println("NOT_YET"); 
  }
}

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

  // --- 3. START I2C S NOVÝMI POISTKAMI ---
  Wire.begin();
  Wire.setClock(50000);

  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);

  // --- Enkodery ---
  pinMode(ENC_LEFT_PIN,  INPUT_PULLUP);
  pinMode(ENC_RIGHT_PIN, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(ENC_LEFT_PIN),  isr_enc_left,  RISING);
  attachInterrupt(digitalPinToInterrupt(ENC_RIGHT_PIN), isr_enc_right, RISING);

  Serial.println("Lasery 1");
  // XSHUT piny su uz OUTPUT a LOW z kroku 2 — len pridame dlhsi reset
  delay(2000);

  // --- 4. POSTUPNÉ ZAPNUTIE LASEROV ---
  Serial.println("Lasery 2");
  init_laser(sensor_back,  XSHUT_B, 0x31, "ZADNY");
  init_laser(sensor_front, XSHUT_F, 0x30, "PREDNY");
  
  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 < 10) {
    err = IMU.init(calib, IMU_ADDRESS);
    delay(200);
    pokus++;
  }
  if (err != 0) {
    Serial.print("Error initializing IMU: ");
    Serial.println(err);
    imu_initialized = false;
  } else {
    Serial.println("IMU initialized");
    IMU.setGyroRange(500);
    measure_YAW_bias(SAMPLE_AMOUNT, 5);
    imu_initialized = true;
  }
  old_time = micros();
  initialized = true;
}

void loop() {
  static char cmd_buf[32];
  static uint8_t cmd_idx = 0;

  while (Serial.available()) {
    char c = Serial.read();
    if (c == '\n' || c == '\r') {
      if (cmd_idx > 0) {
        cmd_buf[cmd_idx] = '\0';

        // Najdi oddelovac ':' a parsuj hodnotu
        char* colon = strchr(cmd_buf, ':');
        int value = 0;
        if (colon) {
          *colon = '\0';
          value = atoi(colon + 1);
        }

        for (int i = 0; i < functions_count; i++) {
          if (strcasecmp(cmd_buf, functions[i].name) == 0) {
            functions[i].function(value);
            break;
          }
        }
        cmd_idx = 0;
      }
    } else if (cmd_idx < sizeof(cmd_buf) - 1) {
      cmd_buf[cmd_idx++] = c;
    }
  }

  if (initialized) {
    update_yaw();
  }
}