#include "toid_behaviors/approach_acorns.hpp"

#include "angles/angles.h"
#include "tf2_geometry_msgs/tf2_geometry_msgs.hpp"

namespace toid
{

ApproachAcorns::ApproachAcorns() : SimpleMove<MoveAction>() {}
ApproachAcorns::~ApproachAcorns() {}

void ApproachAcorns::configureCB()
{
  auto node = node_.lock();
  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".max_vel_accel", rclcpp::ParameterValue(2.0));
  node->get_parameter(behavior_name_ + ".max_vel_accel", max_vel_accel_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".max_vel_decel", rclcpp::ParameterValue(1.0));
  node->get_parameter(behavior_name_ + ".max_vel_decel", max_vel_decel_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".max_vel_speed", rclcpp::ParameterValue(0.4));
  node->get_parameter(behavior_name_ + ".max_vel_speed", max_vel_speed_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".min_vel_speed", rclcpp::ParameterValue(0.10));
  node->get_parameter(behavior_name_ + ".min_vel_speed", min_vel_speed_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".max_angular_speed", rclcpp::ParameterValue(1.0));
  node->get_parameter(behavior_name_ + ".max_angular_speed", max_angular_speed_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".kp", rclcpp::ParameterValue(5.0));
  node->get_parameter(behavior_name_ + ".kp", kp_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".kphi", rclcpp::ParameterValue(5.0));
  node->get_parameter(behavior_name_ + ".kphi", k_phi_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".kdelta", rclcpp::ParameterValue(2.0));
  node->get_parameter(behavior_name_ + ".kdelta", k_delta_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".beta", rclcpp::ParameterValue(0.4));
  node->get_parameter(behavior_name_ + ".beta", beta_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".lambda", rclcpp::ParameterValue(2.0));
  node->get_parameter(behavior_name_ + ".lambda", lambda_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".slow_down_radius", rclcpp::ParameterValue(0.20));
  node->get_parameter(behavior_name_ + ".slow_down_radius", slowdown_radius_);

  nav2_util::declare_parameter_if_not_declared(
    node, behavior_name_ + ".debug_marker", rclcpp::ParameterValue(false));
  node->get_parameter(behavior_name_ + ".debug_marker", debug_marker_);

  target_pose_pub_ = node->create_publisher<visualization_msgs::msg::Marker>("marker", 1);
}

void ApproachAcorns::activateCB()
{
  auto node = node_.lock();
  using namespace std::placeholders;
  acorn_pose_sub_ = node->create_subscription<PoseStamped>(
    "closest_acorn", rclcpp::QoS(2).keep_last(2), std::bind(&ApproachAcorns::acorn_position_cb, this, _1));
  target_pose_pub_->on_activate();
  distance_ = 1.0;
}

void ApproachAcorns::deactivateCB()
{
  acorn_pose_sub_.reset();
  target_pose_pub_->on_deactivate();
}

void ApproachAcorns::acorn_position_cb(PoseStamped::ConstSharedPtr msg)
{
  geometry_msgs::msg::PoseStamped pose_local;
  geometry_msgs::msg::PoseStamped pose_global;
  try {
    pose_local = tf_->transform(*msg, robot_base_frame_);
    pose_global = tf_->transform(*msg, global_frame_);
  } catch (const tf2::TransformException & e) {
    RCLCPP_ERROR_ONCE(logger_, "Transform timeout [showing once]...");
    return;
  }

  double x = pose_local.pose.position.x;
  double y = pose_local.pose.position.y;
  if (x * x + y * y > distance_ + 0.005) {
    return;
  }

  // Smooth out approach
  if (x * x + y * y < 0.42 * 0.42) {
    return;
  }

  double yaw = angles::normalize_angle(tf2::getYaw(pose_global.pose.orientation));

  if (tf2::getYaw(pose_local.pose.orientation) > 0) {
    yaw += M_PI;
  }

  yaw += M_PI / 2;

  pose_global.pose.position.x += std::cos(yaw) * -0.35 - std::sin(yaw) * 0.010;
  pose_global.pose.position.y += std::sin(yaw) * -0.35 + std::cos(yaw) * 0.010;


  std::lock_guard _lock(mutex_);

  const double dx = initial_pose_.position.x - pose_global.pose.position.x;
  const double dy = initial_pose_.position.y - pose_global.pose.position.y;
  double yaw_to_goal = std::atan2(dy,dx);

  tf2::convert(tf2::Quaternion::createFromRPY(0, 0, yaw_to_goal), pose_global.pose.orientation);
  yaw = yaw_to_goal;
  auto [a, b, c] = avg_.push(pose_global.pose);
  distance_ = x * x + y * y;
  new_target_pose_.position.x = a;
  new_target_pose_.position.y = b;
  tf2::convert(tf2::Quaternion::createFromRPY(0, 0, c), new_target_pose_.orientation);
  new_target_angle_ = c;

  if (debug_marker_) {
    visualization_msgs::msg::Marker marker;
    marker.lifetime.sec = 1.0;
    marker.header = pose_global.header;
    marker.pose = new_target_pose_;
    marker.type = visualization_msgs::msg::Marker::ARROW;
    marker.scale.x = 0.02;
    marker.scale.y = 0.02;
    marker.scale.z = 0.02;
    marker.color.a = 1.0;
    marker.color.r = 1.0;
    marker.color.g = 0;
    marker.color.b = 0;
    target_pose_pub_->publish(marker);
  }
}

ResultStatus ApproachAcorns::onStart(
  const std::shared_ptr<const MoveAction::Goal> command, const geometry_msgs::msg::Pose & pose,
  const geometry_msgs::msg::Twist & vel)
{
  std::lock_guard _lock(mutex_);
  distance_ = 1.0;
  new_target_pose_.position.x = command->x;
  new_target_pose_.position.y = command->y;

  backwards_ = command->backwards;
  initial_pose_ = pose;

  new_target_pose_.orientation = tf2::toMsg(tf2::Quaternion::createFromRPY(0, 0, command->theta));
  new_target_angle_ = command->theta;
  target_sign_ = backwards_ ? -1.0 : 1.0;
  max_vel_speed_ = command->max_speed;

  avg_.reset();

  if (command->max_speed == 0) {
    auto node = node_.lock();
    node->get_parameter(behavior_name_ + ".max_vel_speed", max_vel_speed_);
  }

  mode_ = command->mode;

  scl_.k_phi = k_phi_;
  scl_.k_delta = k_delta_;
  scl_.bbeta = beta_;
  scl_.lam = lambda_;
  scl_.slowdown_radius = slowdown_radius_;
  scl_.v_angular_max = max_angular_speed_;
  scl_.v_linear_min = min_vel_speed_;
  scl_.v_linear_max = max_vel_speed_;

  last_speed_ = vel.angular.x;

  return ResultStatus{Status::SUCCEEDED};
}

double ApproachAcorns::distanceToTarget(
  const geometry_msgs::msg::Pose & pose, const geometry_msgs::msg::Point & target_point,
  const double target_theta, bool backwards)
{
  const double dx = target_point.x - pose.position.x;
  const double dy = target_point.y - pose.position.y;
  const double target_sign = backwards ? -1.0 : 1.0;

  return target_sign * (dx * cos(target_theta) + dy * sin(target_theta));
}

double ApproachAcorns::velocityTarget(const double dist_left)
{
  const double lower_bound = last_speed_ - control_duration_ * max_vel_accel_;
  const double upper_bound = last_speed_ + control_duration_ * max_vel_accel_;

  double vel = max_vel_speed_;
  double max_vel_to_stop = 0.8 * std::sqrt(2.0 * max_vel_decel_ * dist_left);
  vel = std::min(vel, max_vel_to_stop);
  return std::clamp(target_sign_ * vel, lower_bound, upper_bound);
}

bool ApproachAcorns::collisionDetection(
  const geometry_msgs::msg::Pose & pose, geometry_msgs::msg::Pose & last_ok_pose)
{
  const int samples = static_cast<int>(0.5 / control_duration_);
  geometry_msgs::msg::Pose current_pose = pose;
  last_ok_pose = pose;
  const bool check_map = !(mode_ & MoveAction::Goal::IGNORE_OBSTACLES);
  for (int i = 0; i < samples; i++) {
    scl_.step(target_pose_, current_pose, control_duration_, backwards_);

    geometry_msgs::msg::Pose2D p;
    p.x = current_pose.position.x;
    p.y = current_pose.position.y;
    p.theta = tf2::getYaw(current_pose.orientation);

    if (check_map && !local_collision_checker_->isCollisionFree(p, i == 0)) {
      return true;
    }

    if (check_rival_collision(p)) {
      return true;
    }

    last_ok_pose = current_pose;
    const double dist_left =
      distanceToTarget(current_pose, target_pose_.position, target_angle_, backwards_);
    if (dist_left < 0.005) {
      return false;
    }
  }
  return false;
}

ResultStatus ApproachAcorns::updateVel(
  const geometry_msgs::msg::Pose & pose, const geometry_msgs::msg::Twist &,
  geometry_msgs::msg::Twist & out_vel)
{
  {
    std::lock_guard _lock(mutex_);
    target_pose_ = new_target_pose_;
    target_angle_ = new_target_angle_;
  }

  double dist_left = distanceToTarget(pose, target_pose_.position, target_angle_, backwards_);

  const double current_yaw = tf2::getYaw(pose.orientation);
  double angle_dist = angles::shortest_angular_distance(current_yaw, target_angle_);

  if (dist_left <= 0.001 && std::fabs(angle_dist) <= 0.030) {
    out_vel.linear.x = 0;
    out_vel.angular.z = 0;
    return ResultStatus{Status::SUCCEEDED};
  }

  geometry_msgs::msg::Pose last_ok_pose;
  if (collisionDetection(pose, last_ok_pose)) {
    dist_left = distanceToTarget(
      pose, last_ok_pose.position, tf2::getYaw(last_ok_pose.orientation), backwards_);
    if (dist_left <= 0.02) {
      out_vel.linear.x = 0;
      out_vel.angular.z = 0;
    } else {
      scl_.v_linear_max = target_sign_ * velocityTarget(dist_left);
      scl_.calculate_vel(last_ok_pose, pose, out_vel, backwards_);
    }

    last_speed_ = out_vel.linear.x;

    RCLCPP_DEBUG_THROTTLE(logger_, *clock_, 1000, "Distance Left: %lf", dist_left);
    RCLCPP_DEBUG_THROTTLE(logger_, *clock_, 1000, "Max Speed: %lf", scl_.v_linear_max);

    return ResultStatus{Status::RUNNING};
  }

  if (dist_left >= 0.001 && dist_left <= 0.02) {
    out_vel.linear.x = velocityTarget(dist_left);
    out_vel.angular.z = std::clamp(kp_ * angle_dist, -max_angular_speed_, max_angular_speed_);
    last_speed_ = out_vel.linear.x;
    return ResultStatus{Status::RUNNING};
  }
  
  if (dist_left <= 0.001) {
    out_vel.linear.x = 0;
    out_vel.angular.z = std::clamp(kp_ * angle_dist, -max_angular_speed_, max_angular_speed_);
    last_speed_ = out_vel.linear.x;
    return ResultStatus{Status::RUNNING};
  }

  scl_.v_linear_max = target_sign_ * velocityTarget(dist_left);
  scl_.calculate_vel(target_pose_, pose, out_vel, backwards_);

  last_speed_ = out_vel.linear.x;

  RCLCPP_DEBUG_THROTTLE(logger_, *clock_, 1000, "Distance Left: %lf", dist_left);
  RCLCPP_DEBUG_THROTTLE(logger_, *clock_, 1000, "Max Speed: %lf", scl_.v_linear_max);

  return ResultStatus{Status::RUNNING};
}

}  // namespace toid

#include "pluginlib/class_list_macros.hpp"
PLUGINLIB_EXPORT_CLASS(toid::ApproachAcorns, nav2_core::Behavior);