#include <mkt_algorithm/diff/diff_go_straight.h>
#include <boost/dll/alias.hpp>
#include <robot/robot.h>

void mkt_algorithm::diff::GoStraight::initialize(
    robot::NodeHandle &nh, const std::string &name, TFListenerPtr tf, robot_costmap_2d::Costmap2DROBOT *costmap_robot, const score_algorithm::TrajectoryGenerator::Ptr &traj)
{
  if (!initialized_)
  {
    nh_ = robot::NodeHandle("~");
    nh_priv_ = robot::NodeHandle(nh, name);
    name_ = name;
    tf_ = tf;
    traj_ = traj;
    costmap_robot_ = costmap_robot;

    this->getParams();
    nh_.param<double>("min_approach_linear_velocity", min_approach_linear_velocity_, 0.1);

    std::vector<robot_geometry_msgs::Point> footprint = costmap_robot_? costmap_robot_->getRobotFootprint() : std::vector<robot_geometry_msgs::Point>();
    if (footprint.size() > 1)
    {
      double min_length = 1e10;
      double max_length = 0.0;
      for (size_t i = 0; i < footprint.size(); ++i)
      {
        const robot_geometry_msgs::Point &p1 = footprint[i];
        const robot_geometry_msgs::Point &p2 = footprint[(i + 1) % footprint.size()]; // wrap around to first point
        double len = std::hypot(p2.x - p1.x, p2.y - p1.y) / 2.0;
        if (len > max_length)
        {
          max_length = len;
        }
        if (len < min_length)
        {
          min_length = len;
        }
      }
      this->min_path_distance_ = min_length > 0.1 ? min_length : 0.1;
      this->max_path_distance_ = max_length > 0.1 ? max_length : 0.1;
    }

    // kalman
    last_actuator_update_ = robot::Time::now();
    n_ = 6; // [x, vx, ax, y, vy, ay, theta, vtheta, atheta]
    m_ = 2; // measurements: x, y, theta
    double dt = control_duration_;

    // Khởi tạo ma trận
    A = Eigen::MatrixXd::Identity(n_, n_);
    C = Eigen::MatrixXd::Zero(m_, n_);
    Q = Eigen::MatrixXd::Zero(n_, n_);
    R = Eigen::MatrixXd::Identity(m_, m_);
    P = Eigen::MatrixXd::Identity(n_, n_);

    for (int i = 0; i < n_; i += 3)
    {
      A(i, i + 1) = dt;
      A(i, i + 2) = 0.5 * dt * dt;
      A(i + 1, i + 2) = dt;
    }

    C(0, 0) = 1;
    C(1, 3) = 1;
    Q(2, 2) = 0.1;
    Q(5, 5) = 0.6;

    R(0, 0) = 0.1;
    R(1, 1) = 0.2;

    P(3, 3) = 0.4;
    P(4, 4) = 0.4;
    P(5, 5) = 0.4;

    kf_ = boost::make_shared<KalmanFilter>(dt, A, C, Q, R, P);
    Eigen::VectorXd x0(n_);
    x0 << 0, 0, 0, 0, 0, 0;
    kf_->init(robot::Time::now().toSec(), x0);

    x_direction_ = y_direction_ = theta_direction_ = 0;
    initialized_ = true;
    robot::log_info("GoStraight Initialized successfully");
  }
}

mkt_msgs::Trajectory2D mkt_algorithm::diff::GoStraight::calculator(
    const robot_nav_2d_msgs::Pose2DStamped &pose, const robot_nav_2d_msgs::Twist2D &velocity)
{
  mkt_msgs::Trajectory2D result;
  result.velocity.x = result.velocity.y = result.velocity.theta = 0.0;
  if (!traj_)
    return result;
  if (compute_plan_.poses.size() < 2)
  {
    robot::log_warning("Local compute plan is available");
    return result;
  }

  robot::Time current_time = robot::Time::now();
  double dt = (current_time - last_actuator_update_).toSec();
  last_actuator_update_ = current_time;

  robot_nav_2d_msgs::Twist2D drive_cmd;
  double sign_x = x_direction_;

  robot_nav_2d_msgs::Twist2D twist;
  traj_->startNewIteration(velocity); // Nhận tốc độ hiện tại từ odom
  robot::Rate r(50);
  while (robot::ok() && traj_->hasMoreTwists())
  {
    twist = traj_->nextTwist();
    // ROS_INFO("Twist: x: %.2f, y: %.2f, theta: %.2f", twist.x, twist.y, twist.theta);
    r.sleep();
  }
  
  drive_cmd.x = sqrt(twist.x * twist.x);
  robot_nav_2d_msgs::Path2D transformed_plan = transform_plan_;
  if (transformed_plan.poses.empty())
  {
    robot::log_warning("Transformed plan is empty. Cannot determine a local plan.");
    return result;
  }

  // Dynamically adjust look ahead distance based on the speed
  double lookahead_dist = this->getLookAheadDistance(twist);
  // Return false if the transformed global plan is empty
  robot_geometry_msgs::Pose2D front = transform_plan_.poses.size() > 3 ? (*(transform_plan_.poses.begin() + 1)).pose : transform_plan_.poses.front().pose;
  lookahead_dist += fabs(front.y);

  // Get lookahead point and publish for visualization
  auto carrot_pose = *getLookAheadPoint(velocity, lookahead_dist, transformed_plan);

  // Carrot distance squared
  const double carrot_dist2 =
      (carrot_pose.pose.x * carrot_pose.pose.x) +
      (carrot_pose.pose.y * carrot_pose.pose.y);

  // Find curvature of circle (k = 1 / R)
  double curvature = 0.0;
  if (carrot_dist2 > 2e-2)
  {
    curvature = 2.0 * carrot_pose.pose.y / carrot_dist2;
  }

  // Constrain linear velocity
  this->applyConstraints(
      0.0, lookahead_dist, curvature, twist,
      this->costAtPose(pose.pose.x, pose.pose.y),
      drive_cmd.x, sign_x);


  if (std::hypot(carrot_pose.pose.x, carrot_pose.pose.y) > min_journey_squared_)
  {
    robot_nav_2d_msgs::Pose2DStamped end = transform_plan_.poses.back();
    robot_nav_2d_msgs::Pose2DStamped start = transform_plan_.poses[transform_plan_.poses.size() - 2];
    double dx = end.pose.x - start.pose.x;
    double dy = end.pose.y - start.pose.y;
    drive_cmd.theta = atan2(dy, dx);

    //[-π/2, π/2]
    if (drive_cmd.theta > M_PI / 2)
      drive_cmd.theta -= M_PI;
    else if (drive_cmd.theta < -M_PI / 2)
      drive_cmd.theta += M_PI;

    drive_cmd.theta = std::clamp(drive_cmd.theta, -min_vel_theta_, min_vel_theta_);
  }
  else
    drive_cmd.theta = 0.0;

  // populate and return message
  if (fabs(carrot_pose.pose.x) > 1e-9 || carrot_pose.pose.y > 1e-9)
    drive_cmd.x *= fabs(cos(std::atan2(carrot_pose.pose.y, carrot_pose.pose.x)));
    
  double v_max = sign_x > 0 ? traj_->getTwistLinear(true).x : traj_->getTwistLinear(false).x;
  const double vel_x_reduce = std::min(fabs(v_max), 0.3);
  double journey_plan = compute_plan_.poses.empty() ? 0 : journey(compute_plan_.poses, 0, compute_plan_.poses.size() - 1);
  const double scale = fabs(velocity.x) * lookahead_time_ * 0.9;
  const double min_S = min_lookahead_dist_ + max_path_distance_ + scale, max_S = max_lookahead_dist_ + max_path_distance_ + scale;
  double d_reduce = std::clamp(journey_plan, min_S, max_S);
  double d_begin_reduce = std::clamp(d_reduce * 0.2, 0.4, 1.0);
  double cosine_factor_begin_reduce = 0.5 * (1.0 + cos(M_PI * (1.0 - fabs(journey_plan) / d_begin_reduce)));
  double v_min = 
    journey_plan > d_begin_reduce ? vel_x_reduce : (vel_x_reduce - min_approach_linear_velocity_) * cosine_factor_begin_reduce + min_approach_linear_velocity_;
  v_min *= sign_x;
  
  double effective_journey = getEffectiveDistance(carrot_pose, journey_plan);
  double decel_factor = computeDecelerationFactor(effective_journey, d_reduce);
  journey_plan += max_path_distance_;
  drive_cmd.x = journey_plan >= d_reduce ? drive_cmd.x : (drive_cmd.x - v_min) * decel_factor + v_min;

  // drive_cmd.x = journey_plan > d_reduce ? drive_cmd.x : (drive_cmd.x - v_min) * sin((journey_plan / d_reduce) * (M_PI / 2)) + v_min;

  Eigen::VectorXd y(2);
  y << drive_cmd.x, drive_cmd.theta;

  // Cập nhật lại A nếu dt thay đổi
  for (int i = 0; i < n_; ++i)
    for (int j = 0; j < n_; ++j)
      A(i, j) = (i == j ? 1.0 : 0.0);
  for (int i = 0; i < n_; i += 3)
  {
    A(i, i + 1) = dt;
    A(i, i + 2) = 0.5 * dt * dt;
    A(i + 1, i + 2) = dt;
  }
  kf_->update(y, dt, A);

  if (drive_cmd.x != v_max)
  {
    double min = drive_cmd.x < v_max ? drive_cmd.x : v_max;
    double max = drive_cmd.x > v_max ? drive_cmd.x : v_max;
    drive_cmd.x = std::clamp(kf_->state()[0], min, max);
  }
  else
  {
    drive_cmd.x = std::clamp(kf_->state()[0], v_max, v_max);
  }

  drive_cmd.x = fabs(drive_cmd.x) >= min_approach_linear_velocity_ ? drive_cmd.x : std::copysign(min_approach_linear_velocity_, sign_x);
  drive_cmd.theta = kf_->state()[3];
  result.velocity = drive_cmd;
  return result;
}

score_algorithm::ScoreAlgorithm::Ptr mkt_algorithm::diff::GoStraight::create()
{
  return std::make_shared<mkt_algorithm::diff::GoStraight>();
}

BOOST_DLL_ALIAS(mkt_algorithm::diff::GoStraight::create, MKTAlgorithmDiffGoStraight)