#include <algorithm>

#include <spdlog/spdlog.h>
#include <spdlog/sinks/stdout_color_sinks.h>

#include "include/ClosedLoopMotorController.hpp"
#include "include/Motors.hpp"
#include "include/Encoders.hpp"

#define MS_IN_S 1000
#define US_IN_S (1000 * MS_IN_S)

void ClosedLoopMotorController::set_power(uint8_t port, double power) {
  std::lock_guard<std::recursive_mutex> lock(speed_control_mutex);
  control_modes.at(port) = ControlMode::POWER;
  Motors::set_power(port, power);
}

void ClosedLoopMotorController::set_speed(uint8_t port, double speed_ms) {
  std::lock_guard<std::recursive_mutex> lock(speed_control_mutex);
  speed_targets.at(port) = speed_ms;
  if (control_modes.at(port) != ControlMode::SPEED) {
    pids.at(port) = PID(
        ROBOT_MOTOR_SPEED_CONTROL_KP, ROBOT_MOTOR_SPEED_CONTROL_KI, ROBOT_MOTOR_SPEED_CONTROL_KD,
        ROBOT_MOTOR_SPEED_CONTROL_RAMP, 100.0);
    control_modes.at(port) = ControlMode::SPEED;
  }
}

ClosedLoopMotorController::ClosedLoopMotorController() {
  speed_control_thread = std::thread(&ClosedLoopMotorController::speed_control_loop, this);
  speed_control_thread.detach();
}

[[noreturn]] void ClosedLoopMotorController::speed_control_loop() {
  auto sleep_duration = std::chrono::microseconds(US_IN_S / ROBOT_MOTOR_SPEED_CONTROL_RATE_HZ);
  while (true) {
    std::this_thread::sleep_for(sleep_duration);

    std::lock_guard<std::recursive_mutex> lock(speed_control_mutex);
    bool should_control_speed = std::find(control_modes.begin(), control_modes.end(), ControlMode::SPEED) != control_modes.end();
    if (!should_control_speed) {
      continue;
    }

    auto speeds_rad_s = Encoders::getInstance().get_velocities_rad_s();
    for (int i = 0; i < ROBOT_MOTOR_COUNT; i++) {
      if (control_modes.at(i) == ControlMode::SPEED) {
        spdlog::debug("CLMC: {}", i);
        double power = pids.at(i)(speed_targets.at(i), speeds_rad_s.at(i));
        Motors::set_power(i, power);
      }
    }
  }
}

void ClosedLoopMotorController::generate_step_response_data() {
  auto console = spdlog::stdout_color_mt("generate_step_response_data");
  console->set_pattern("%v");

  auto input = 30.0;

  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  pids.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = PID(
      1.0, 0.0, 0.0,
      1000.0, 100.0);
  pids.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = PID(
      1.0, 0.0, 0.0,
      1000.0, 100.0);
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = ControlMode::SPEED;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = ControlMode::SPEED;

  std::this_thread::sleep_for(std::chrono::seconds(1));

  auto start_time = std::chrono::steady_clock::now();
  for (int i = 0; i < 1000; ++i) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    double delta_seconds = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(std::chrono::steady_clock::now() - start_time).count() / US_IN_S;
    auto output = Encoders::getInstance().get_velocities_rad_s().at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT);
    console->info("{:03.4f}; {:03.4f}; {:03.4f}", delta_seconds, input, output);
  }

  input = 40.0;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  for (int i = 0; i < 1000; ++i) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    double delta_seconds = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(std::chrono::steady_clock::now() - start_time).count() / US_IN_S;
    auto output = Encoders::getInstance().get_velocities_rad_s().at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT);
    console->info("{:03.4f}; {:03.4f}; {:03.4f}", delta_seconds, input, output);
  }

  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = 0;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = 0;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = ControlMode::NONE;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = ControlMode::NONE;
}

void ClosedLoopMotorController::generate_tuned_step_response_data() {
  auto console = spdlog::stdout_color_mt("generate_tuned_step_response_data");
  console->set_pattern("%v");

  auto input = 2 * M_PI;

  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  pids.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = PID(
      5.5, 50.0, 0.0,
      4 * M_PI, 100.0);
  pids.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = PID(
      5.5, 50.0, 0.0,
      4 * M_PI, 100.0);
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = ControlMode::SPEED;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = ControlMode::SPEED;

  std::this_thread::sleep_for(std::chrono::seconds(2));

  auto start_time = std::chrono::steady_clock::now();
  for (int i = 0; i < 1500; ++i) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    double delta_seconds = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(std::chrono::steady_clock::now() - start_time).count() / US_IN_S;
    auto output = Encoders::getInstance().get_velocities_rad_s().at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT);
    console->info("{:03.4f}; {:03.4f}; {:03.4f}", delta_seconds, input, output);
  }

  input = 4 * M_PI;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  for (int i = 0; i < 1500; ++i) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    double delta_seconds = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(std::chrono::steady_clock::now() - start_time).count() / US_IN_S;
    auto output = Encoders::getInstance().get_velocities_rad_s().at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT);
    console->info("{:03.4f}; {:03.4f}; {:03.4f}", delta_seconds, input, output);
  }

  input = M_PI;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  for (int i = 0; i < 1500; ++i) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    double delta_seconds = std::chrono::duration_cast<std::chrono::duration<double, std::micro>>(std::chrono::steady_clock::now() - start_time).count() / US_IN_S;
    auto output = Encoders::getInstance().get_velocities_rad_s().at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT);
    console->info("{:03.4f}; {:03.4f}; {:03.4f}", delta_seconds, input, output);

    input += 0.01;
    speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_LEFT_MULT;
    speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = input * ROBOT_ODOMETRY_CONTROLLER_RIGHT_MULT;
  }

  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = 0;
  speed_targets.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = 0;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_LEFT_PORT) = ControlMode::NONE;
  control_modes.at(ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT) = ControlMode::NONE;
}

void ClosedLoopMotorController::calibrate_wheel_ticks(double turns, double ticks_per_turn) {
  auto k_p = 0.0001;
  auto power = 25.0;
  auto min_power = 7.5;
  auto port = ROBOT_ODOMETRY_CONTROLLER_RIGHT_PORT;

  auto ticks = 0.0;
  auto target_ticks = ticks_per_turn * turns;
  while (ticks != target_ticks) {
    auto power_mult = fmax(-1.0, fmin(1.0, (target_ticks - ticks) * k_p));
    auto adjusted_power = power * power_mult;
    if (adjusted_power < 0 && adjusted_power > -min_power) {
      adjusted_power = -min_power;
    } else if (adjusted_power > 0 && adjusted_power < min_power) {
      adjusted_power = min_power;
    }

    set_power(port, adjusted_power);

    spdlog::info("{:03.4f}; {:03.4f}; {:03.4f};",
                 ticks, target_ticks, adjusted_power);

    std::this_thread::sleep_for(std::chrono::milliseconds(2));
    ticks = Encoders::getInstance().get_positions().at(port);
  }

  set_power(port, 0);
  std::this_thread::sleep_for(std::chrono::milliseconds(100));

}
