#include <robot/robot.h>
#include <score_algorithm/score_algorithm.h>
#include <robot_nav_2d_utils/path_ops.h>
#include <robot_nav_2d_utils/conversions.h>
#include <angles/angles.h>

unsigned int score_algorithm::ScoreAlgorithm::getGoalIndex(const robot_geometry_msgs::Pose2D &robot_pose, const std::vector<robot_nav_2d_msgs::Pose2DStamped> &plan,
                                                           const double distance, unsigned int start_index, unsigned int last_valid_index) const
{
  if (start_index >= last_valid_index)
    return last_valid_index;

  unsigned int goal_index = start_index + 1;
  const double start_direction_x = plan[goal_index].pose.x - plan[start_index].pose.x;
  const double start_direction_y = plan[goal_index].pose.y - plan[start_index].pose.y;
  if (fabs(start_direction_x) > 1e-9 || fabs(start_direction_y) > 1e-9)
  { // make sure that atan2 is defined
    double start_angle = atan2(start_direction_y, start_direction_x);
    double S = sqrt(pow(start_direction_x, 2) + pow(start_direction_y, 2));
    // double angle_threshold = 0.0;
    for (unsigned int i = start_index + 1; i <= last_valid_index; i++)
    {
      const double current_direction_x = plan[i].pose.x - plan[i - 1].pose.x;
      const double current_direction_y = plan[i].pose.y - plan[i - 1].pose.y;
      const double dd = sqrt(pow(current_direction_x, 2) + pow(current_direction_y, 2));
      S += dd;
      if (fabs(current_direction_x) > 1e-9 || fabs(current_direction_y) > 1e-9)
      {
        double current_angle = atan2(current_direction_y, current_direction_x);
        goal_index = i;
        if (fabs(remainder(start_angle - current_angle, 2 * M_PI)) >= angle_threshold_ || S >= distance)
        {
          break;
        }
      }
    }
  }
  if (last_valid_index - goal_index == 1)
    goal_index = last_valid_index;
  return goal_index;
}

bool score_algorithm::ScoreAlgorithm::findSubGoalIndex(const std::vector<robot_nav_2d_msgs::Pose2DStamped> &plan, std::vector<unsigned int> &seq_s, std::vector<robot_nav_2d_msgs::Pose2DStamped> &mutes)
{
  if (plan.empty())
  {
    return false;
  }
  double x_direction_saved = 0.0;
  robot_geometry_msgs::PoseArray poseArrayMsg;
  poseArrayMsg.header.stamp = robot::Time::now();
  poseArrayMsg.header.frame_id = costmap_robot_->getGlobalFrameID();

  robot_nav_2d_msgs::Pose2DStamped p1, p2, p3;

  if (plan.size() < 3)
  {
    seq_s.push_back(plan.back().header.seq);
    mutes.push_back(plan.back());
  }
  else
  {
    unsigned int i = 0;
    p1 = plan[i];
    for (i = 0; i < (unsigned int)plan.size() - 2; i++)
    {
      double egde1_angle, egde2_angle;
      p2 = plan[i + 1];
      if ((fabs(p2.pose.y - p1.pose.y) > 1e-9 || fabs(p2.pose.x - p1.pose.x) > 1e-9))
      {
        egde1_angle = atan2(p2.pose.y - p1.pose.y, p2.pose.x - p1.pose.x);
      }
      else
        continue;

      const double target_direction = cos(fabs(egde1_angle - p2.pose.theta));

      if (target_direction * x_direction_saved < 0.0)
      {
        if (std::find(seq_s.begin(), seq_s.end(), plan[i].header.seq) == seq_s.end())
        {
          seq_s.push_back(plan[i].header.seq);
          mutes.push_back(plan[i]);
          robot_geometry_msgs::Pose pose = robot_nav_2d_utils::pose2DToPoseStamped(plan[i]).pose;
          poseArrayMsg.poses.push_back(pose);
        }
      }
      x_direction_saved = target_direction;

      p3 = plan[i + 2];
      if ((fabs(p3.pose.y - p2.pose.y) > 1e-9 || fabs(p3.pose.x - p2.pose.x) > 1e-9))
      {
        egde2_angle = atan2(p3.pose.y - p2.pose.y, p3.pose.x - p2.pose.x);
        p1 = p2;
      }
      else
        continue;

      double thresh_angle = 0.3;
      double angle = fabs(remainder(egde1_angle - egde2_angle, 2 * M_PI));
      if (angle >= thresh_angle || angle >= M_PI - thresh_angle)
      {
        if (std::find(seq_s.begin(), seq_s.end(), plan[i].header.seq) == seq_s.end())
        {
          seq_s.push_back(plan[i + 1].header.seq);
          mutes.push_back(plan[i + 1]);
          robot_geometry_msgs::Pose pose = robot_nav_2d_utils::pose2DToPoseStamped(plan[i + 1]).pose;
          poseArrayMsg.poses.push_back(pose);
        }
        p2 = p3;
      }
    }
  }

  if (seq_s.empty())
  {
    seq_s.push_back(plan.back().header.seq);
    mutes.push_back(plan.back());
    robot_geometry_msgs::Pose pose = robot_nav_2d_utils::pose2DToPoseStamped(plan.back()).pose;
    poseArrayMsg.poses.push_back(pose);
  }
  else
  {
    if (seq_s.back() < plan.back().header.seq)
    {
      seq_s.push_back(plan.back().header.seq);
      mutes.push_back(plan.back());
      robot_geometry_msgs::Pose pose = robot_nav_2d_utils::pose2DToPoseStamped(plan.back()).pose;
      poseArrayMsg.poses.push_back(pose);
    }
  }
  return true;
}

bool score_algorithm::ScoreAlgorithm::computePlanCommand(const robot_nav_2d_msgs::Pose2DStamped &robot_pose, const robot_nav_2d_msgs::Twist2D &velocity, const double &S,
                                                         const robot_nav_2d_msgs::Path2D &global_plan, unsigned int &start_index_global_plan, unsigned int &goal_index_global_plan, robot_nav_2d_msgs::Path2D &result)
{
  result.poses.clear();
  if (global_plan.poses.empty())
  {
    robot::log_error("ScoreAlgorithm: The global plan was empty.");
    return false;
  }

  std::vector<unsigned int> seq_s;
  std::vector<robot_nav_2d_msgs::Pose2DStamped> mutes;
  if (!this->findSubGoalIndex(global_plan.poses, seq_s, mutes))
  {
    robot::log_error("ScoreAlgorithm: Cannot find SubGoal Index");
    return false;
  }

  unsigned int sub_goal_index;
  unsigned int sub_start_index;

  bool found_sub_goal_index_saved = false, found_sub_start_index_saved = false;
  for (unsigned int i = 0; i < (unsigned int)global_plan.poses.size(); ++i)
  {
    if (global_plan.poses[i].header.seq == sub_goal_seq_saved_)
    {
      sub_goal_index_saved_ = i;
      found_sub_goal_index_saved = true;
    }
    if (global_plan.poses[i].header.seq == sub_start_seq_saved_)
    {
      sub_start_index_saved_ = i;
      found_sub_start_index_saved = true;
    }
    if (found_sub_goal_index_saved && found_sub_start_index_saved)
    {
      break;
    }
  }

  sub_goal_index = !found_sub_goal_index_saved ? (unsigned int)global_plan.poses.size() : sub_goal_index_saved_;
  sub_start_index = !found_sub_start_index_saved ? 0 : sub_start_index_saved_;

  std::vector<unsigned int> index_s;
  for (auto seq : seq_s)
  {
    for (unsigned int i = 0; i < (unsigned int)global_plan.poses.size(); ++i)
    {
      if (global_plan.poses[i].header.seq == seq)
      {
        index_s.push_back(i);
        break;
      }
    }
  }

  if (index_s.empty())
  {
    sub_goal_index = (unsigned int)global_plan.poses.size();
    sub_goal_seq_saved_ = global_plan.poses.back().header.seq;
  }
  else
  {
    bool is_stopped = stopped(velocity, rot_stopped_velocity_, trans_stopped_velocity_);
    if (is_stopped && index_s.front() > sub_goal_index_saved_ && index_s.front() < (unsigned int)global_plan.poses.size())
    {
      // sub_goal_index = (unsigned int)index_s.front() > 0 ? (unsigned int)index_s.front() - 1 : (unsigned int)index_s.front();
      sub_goal_index = index_s.front();
      sub_goal_seq_saved_ = global_plan.poses[sub_goal_index].header.seq;
    }

    // Process index_s with multiple elements
    if (index_s.size() > 1)
    {
      for (size_t i = 1; i < index_s.size(); ++i)
      {
        if (index_s[i - 1] >= (unsigned int)global_plan.poses.size() || index_s[i] >= (unsigned int)global_plan.poses.size())
        {
          continue;
        }

        double dx = robot_pose.pose.x - global_plan.poses[index_s[i - 1]].pose.x;
        double dy = robot_pose.pose.y - global_plan.poses[index_s[i - 1]].pose.y;
        double theta = angles::normalize_angle(robot_pose.pose.theta - global_plan.poses[index_s[i - 1]].pose.theta);
        double x = cos(global_plan.poses[index_s[i - 1]].pose.theta) * dx + sin(global_plan.poses[index_s[i - 1]].pose.theta) * dy;
        double y = -sin(global_plan.poses[index_s[i - 1]].pose.theta) * dx + cos(global_plan.poses[index_s[i - 1]].pose.theta) * dy;

        if (std::abs(std::sqrt(dx * dx + dy * dy)) <= xy_local_goal_tolerance_ + 0.2)
        {
          double tolerance = fabs(cos(theta)) >= fabs(sin(theta)) ? x : y;
          if (fabs(tolerance) <= xy_local_goal_tolerance_)
          {

            if (index_s[i] > sub_goal_index_saved_)
            {
              sub_goal_index = (i < index_s.size() - 1) ? index_s[i] : (unsigned int)global_plan.poses.size() - 1;
              sub_goal_seq_saved_ = global_plan.poses[sub_goal_index].header.seq;
              old_xy_goal_tolerance_ = 0.0;
              if (follow_step_path_)
              {
                sub_start_index = index_s[i - 1];
                sub_start_seq_saved_ = global_plan.poses[sub_start_index].header.seq;
              }
              break;
            }
          }
        }
        else
          old_xy_goal_tolerance_ = fabs(cos(theta)) >= fabs(sin(theta)) ? x : y;
      }
    }
    else
    {
      // Validate sub_goal_index
      if (sub_goal_index >= (unsigned int)global_plan.poses.size() - 1)
      {
        sub_goal_index = (unsigned int)global_plan.poses.size() - 1;
        sub_goal_seq_saved_ = global_plan.poses[sub_goal_index].header.seq;
      }
      else
      {
        double dx = robot_pose.pose.x - global_plan.poses[sub_goal_index].pose.x;
        double dy = robot_pose.pose.y - global_plan.poses[sub_goal_index].pose.y;
        double theta = angles::normalize_angle(robot_pose.pose.theta - global_plan.poses[sub_goal_index].pose.theta);
        double x = cos(global_plan.poses[sub_goal_index].pose.theta) * dx + sin(global_plan.poses[sub_goal_index].pose.theta) * dy;
        double y = -sin(global_plan.poses[sub_goal_index].pose.theta) * dx + cos(global_plan.poses[sub_goal_index].pose.theta) * dy;
        if (std::abs(std::sqrt(dx * dx + dy * dy)) <= xy_local_goal_tolerance_ + 0.2)
        {
          double tolerance = fabs(cos(theta)) >= fabs(sin(theta)) ? x : y;
          if (fabs(tolerance) <= xy_local_goal_tolerance_)
          {
            sub_goal_index = (unsigned int)global_plan.poses.size() - 1;
            sub_goal_seq_saved_ = global_plan.poses[sub_goal_index].header.seq;
            old_xy_goal_tolerance_ = 0.0;
            robot::log_info("ScoreAlgorithm: Single sub-goal reached, setting sub_goal_index to %d", sub_goal_index);
          }
        }
        else
          old_xy_goal_tolerance_ = fabs(cos(theta)) >= fabs(sin(theta)) ? x : y;
      }
    }
  }
  if (sub_start_index >= sub_goal_index)
  {
    robot::log_error("ScoreAlgorithm: Sub path is empty");
    return false;
  }

  // find the "start pose", i.e. the pose on the plan closest to the robot that is also on the local map
  // Tìm điểm pose bắt đầu trên kế hoạch đường đi gần nhất với robot đồng thời có tồn tại trong local map
  unsigned int start_index = 0; // số vị trí pose trong kế hoạch đường đi
  unsigned int closet_index = 0;
  double distance_to_start = std::numeric_limits<double>::infinity(); // Xét giá trị là vô cùng oo
  bool started_path = false;
  for (unsigned int i = sub_start_index; i < sub_goal_index; ++i)
  {
    // Still on the costmap. Continue.
    double distance = fabs(robot_nav_2d_utils::poseDistance(global_plan.poses[i].pose, robot_pose.pose)); // Tính khoảng cách từ vị trí của robot đến pose đang xét
    if (distance_to_start > distance)
    {
      start_index = i;
      closet_index = start_index;
      distance_to_start = distance;
      started_path = true;
    }
  }

  // Nếu khồng tìm được điểm gần với robot nhất thì trả về false
  if (!started_path)
  {
    robot::log_error("ScoreAlgorithm: None of the points of the global plan were in the local costmap.");
    return false;
  }
  std::stringstream ss;
  start_index_saved_vt_.push_back(global_plan.poses[start_index].header.seq);
  if (start_index_saved_vt_.size() > index_samples_)
    start_index_saved_vt_.erase(start_index_saved_vt_.begin());

  unsigned num = 0;
  if (!start_index_saved_vt_.empty())
  {
    for (int i = 0; i < (int)start_index_saved_vt_.size() - 1; i++)
    {
      if (start_index_saved_vt_[i] == start_index_saved_vt_[i + 1])
        num++;
      ss << start_index_saved_vt_[i] << " ";
    }
  }
  if (num < index_samples_ - 1 && !start_index_saved_vt_.empty())
  {
    auto max_iter = *(std::max_element(start_index_saved_vt_.begin(), start_index_saved_vt_.end()));
    unsigned int max_start_index = start_index;
    for (unsigned int i = 0; i < (unsigned int)global_plan.poses.size(); i++)
    {
      if (global_plan.poses[i].header.seq == max_iter)
      {
        max_start_index = i;
        break;
      }
    }
    double S1 = journey(global_plan.poses, start_index, (unsigned int)global_plan.poses.size() - 1);
    double S2 = journey(global_plan.poses, max_start_index, (unsigned int)global_plan.poses.size() - 1);

    if (S2 - S1)
    {
      start_index = max_start_index;
    }
  }
  // ROS_INFO_STREAM_THROTTLE(0.1, ss.str() << "\n");

  // find the "last valid pose", i.e. the last pose on the plan after the start pose that is still on the local map
  // Tìm điểm pose cuôi cùng sau điểm pose bắt đầu và có tồn tại trong local map
  if (start_index < 1)
    start_index = 1;
  // unsigned int last_valid_index = (unsigned int)global_plan.poses.size() - 1;
  unsigned int last_valid_index = sub_goal_index;

  // find the "goal pose" by walking along the plan as long as each step leads far enough away from the starting point,
  // i.e. is within angle_threshold_ of the starting direction.
  // Tìm điểm mục tiêu chạy từ điểm start pose đến last valid pose dọc thep plan
  unsigned int goal_index = start_index;
  while (goal_index < last_valid_index)
  {
    // Tìm kiếm vị trí điểm mục tiêu phù hợp
    goal_index = getGoalIndex(robot_pose.pose, global_plan.poses, S, start_index, last_valid_index);

    // check if goal already reached
    bool goal_already_reached = false;

    for (auto &reached_intermediate_goal : reached_intermediate_goals_)
    {
      double distance = fabs(robot_nav_2d_utils::poseDistance(reached_intermediate_goal.pose, global_plan.poses[goal_index].pose));
      if (distance < xy_local_goal_tolerance_)
      {
        goal_already_reached = true;
        // choose a new start_index by walking along the plan until we're outside xy_local_goal_tolerance_ and try again
        // (start_index is now > goal_index)
        for (start_index = goal_index; start_index < last_valid_index; ++start_index)
        {
          distance = robot_nav_2d_utils::poseDistance(reached_intermediate_goal.pose, global_plan.poses[start_index].pose);
          if (distance >= xy_local_goal_tolerance_)
            break;
        }
        break;
      }
    }

    if (!goal_already_reached)
    {
      // new goal not in reached_intermediate_goals_
      double distance = fabs(robot_nav_2d_utils::poseDistance(global_plan.poses[start_index].pose, global_plan.poses[goal_index].pose));
      if (distance < xy_local_goal_tolerance_)
      {
        // the robot has currently reached the goal
        reached_intermediate_goals_.push_back(global_plan.poses[goal_index]);
        robot::log_info("ScoreAlgorithm: Number of reached_intermediate goals: %d", reached_intermediate_goals_.size());
      }
      else
      {
        // we've found a new goal!
        break;
      }
    }
  }

  if (start_index > goal_index)
    start_index = goal_index;
  if (goal_index > last_valid_index)
    goal_index = last_valid_index;
  if ((unsigned int)global_plan.poses.size() == 2 && !global_plan.poses.empty())
  {
    start_index = 0;
    goal_index = 1;
  }

  if (goal_index == (unsigned int)global_plan.poses.size() - 1)
    start_index = (unsigned int)global_plan.poses.size() - 2;

  start_index_global_plan = start_index;
  goal_index_global_plan = goal_index;

  robot_nav_2d_msgs::Pose2DStamped sub_pose;
  sub_pose = global_plan.poses[closet_index];

#ifdef BUILD_WITH_ROS
  {
    robot_geometry_msgs::PoseStamped sub_pose_stamped = robot_nav_2d_utils::pose2DToPoseStamped(sub_pose);
    geometry_msgs::PoseStamped sub_pose_stamped_ros;
    sub_pose_stamped_ros.header.stamp = ros::Time::now();
    sub_pose_stamped_ros.header.frame_id = sub_pose_stamped.header.frame_id;
    sub_pose_stamped_ros.pose.position.x = sub_pose_stamped.pose.position.x;
    sub_pose_stamped_ros.pose.position.y = sub_pose_stamped.pose.position.y;
    sub_pose_stamped_ros.pose.position.z = 0.9;
    sub_pose_stamped_ros.pose.orientation.x = sub_pose_stamped.pose.orientation.x;
    sub_pose_stamped_ros.pose.orientation.y = sub_pose_stamped.pose.orientation.y;
    sub_pose_stamped_ros.pose.orientation.z = sub_pose_stamped.pose.orientation.z;
    sub_pose_stamped_ros.pose.orientation.w = sub_pose_stamped.pose.orientation.w;
    closet_robot_goal_pub_.publish(sub_pose_stamped_ros);
  }
#endif

  robot_nav_2d_msgs::Pose2DStamped sub_goal;
  sub_goal = global_plan.poses[goal_index];
#ifdef BUILD_WITH_ROS
  {
    robot_geometry_msgs::PoseStamped sub_goal_stamped = robot_nav_2d_utils::pose2DToPoseStamped(sub_goal);
    geometry_msgs::PoseStamped sub_goal_stamped_ros;
    sub_goal_stamped_ros.header.stamp = ros::Time::now();
    sub_goal_stamped_ros.header.frame_id = sub_goal_stamped.header.frame_id;
    sub_goal_stamped_ros.pose.position.x = sub_goal_stamped.pose.position.x;
    sub_goal_stamped_ros.pose.position.y = sub_goal_stamped.pose.position.y;
    sub_goal_stamped_ros.pose.position.z = 0.9;
    sub_goal_stamped_ros.pose.orientation.x = sub_goal_stamped.pose.orientation.x;
    sub_goal_stamped_ros.pose.orientation.y = sub_goal_stamped.pose.orientation.y;
    sub_goal_stamped_ros.pose.orientation.z = sub_goal_stamped.pose.orientation.z;
    sub_goal_stamped_ros.pose.orientation.w = sub_goal_stamped.pose.orientation.w;
    current_goal_pub_.publish(sub_goal_stamped_ros);
  }
#endif

  result.header = global_plan.header;
  for (int i = closet_index; i <= goal_index; i++)
  {
    result.poses.push_back(global_plan.poses[i]);
  }
  return true;
}

double score_algorithm::ScoreAlgorithm::journey(const std::vector<robot_nav_2d_msgs::Pose2DStamped> &plan, const unsigned int start_index, const unsigned int last_valid_index) const
{
  double S = 0;
  if (last_valid_index - start_index > 1)
  {
    for (int i = start_index; i < last_valid_index; i++)
    {
      S += robot_nav_2d_utils::poseDistance(plan[i].pose, plan[i + 1].pose);
    }
  }
  return S;
}

bool score_algorithm::ScoreAlgorithm::stopped(const robot_nav_2d_msgs::Twist2D &velocity, const double &rot_stopped_velocity, const double &trans_stopped_velocity)
{
  return fabs(velocity.theta) <= rot_stopped_velocity && fabs(velocity.x) <= trans_stopped_velocity && fabs(velocity.y) <= trans_stopped_velocity;
}

void score_algorithm::ScoreAlgorithm::clear()
{
  reached_intermediate_goals_.clear();
  start_index_saved_vt_.clear();
  sub_goal_index_saved_ = sub_goal_seq_saved_ = 0;
  sub_start_index_saved_, sub_start_seq_saved_ = 0;
}