#include <ros/ros.h>

#include <nav_core/base_global_planner.h>
#include <costmap_2d/costmap_2d_ros.h>
#include <costmap_2d/cost_values.h>
#include <geometry_msgs/PoseStamped.h>
#include <nav_msgs/Path.h>
#include <pluginlib/class_list_macros.h>

#include <algorithm>
#include <cmath>
#include <limits>
#include <random>
#include <string>
#include <vector>

namespace rrt_star_global_planner
{

struct Point2D
{
  double x;
  double y;
};

struct TreeNode
{
  Point2D point;
  int parent;
  double cost;
};

class RRTStarGlobalPlanner : public nav_core::BaseGlobalPlanner
{
public:
  RRTStarGlobalPlanner()
    : initialized_(false),
      costmap_ros_(nullptr),
      costmap_(nullptr),
      rng_(std::random_device{}())
  {
  }

  RRTStarGlobalPlanner(std::string name, costmap_2d::Costmap2DROS* costmap_ros)
    : initialized_(false),
      costmap_ros_(nullptr),
      costmap_(nullptr),
      rng_(std::random_device{}())
  {
    initialize(name, costmap_ros);
  }

  void initialize(std::string name, costmap_2d::Costmap2DROS* costmap_ros) override
  {
    if (initialized_)
    {
      ROS_WARN("RRTStarGlobalPlanner has already been initialized.");
      return;
    }

    if (costmap_ros == nullptr)
    {
      ROS_ERROR("RRTStarGlobalPlanner received null costmap_ros.");
      return;
    }

    costmap_ros_ = costmap_ros;
    costmap_ = costmap_ros_->getCostmap();
    global_frame_ = costmap_ros_->getGlobalFrameID();

    ros::NodeHandle private_nh("~/" + name);
    ros::NodeHandle nh;

    plan_pub_ = nh.advertise<nav_msgs::Path>("/move_base/RRTStarGlobalPlanner/plan", 1, true);

    private_nh.param("max_iter", max_iter_, 3000);
    private_nh.param("step_size", step_size_, 0.30);
    private_nh.param("search_radius", search_radius_, 0.80);
    private_nh.param("goal_radius", goal_radius_, 0.50);
    private_nh.param("normal_width_ratio", normal_width_ratio_, 0.45);
    private_nh.param("max_sample_retry", max_sample_retry_, 1000);
    private_nh.param("d_dichotomy", d_dichotomy_, 0.20);
    private_nh.param("stop_at_first_path", stop_at_first_path_, true);
    private_nh.param("allow_unknown", allow_unknown_, false);
    private_nh.param("collision_check_step", collision_check_step_, 0.05);

    initialized_ = true;

    ROS_INFO("RRTStarGlobalPlanner initialized.");
    ROS_INFO("global_frame = %s", global_frame_.c_str());
    ROS_INFO("full path topic = /move_base/RRTStarGlobalPlanner/plan");
  }

  bool makePlan(const geometry_msgs::PoseStamped& start,
                const geometry_msgs::PoseStamped& goal,
                std::vector<geometry_msgs::PoseStamped>& plan) override
  {
    plan.clear();

    if (!initialized_)
    {
      ROS_ERROR("RRTStarGlobalPlanner has not been initialized.");
      return false;
    }

    Point2D start_pt{start.pose.position.x, start.pose.position.y};
    Point2D goal_pt{goal.pose.position.x, goal.pose.position.y};

    if (!isPointFree(start_pt))
    {
      ROS_WARN("RRTStarGlobalPlanner: start point is not free.");
      return false;
    }

    if (!isPointFree(goal_pt))
    {
      ROS_WARN("RRTStarGlobalPlanner: goal point is not free.");
      return false;
    }

    std::vector<Point2D> path_points;
    bool success = planRRTStar(start_pt, goal_pt, path_points);

    if (!success || path_points.empty())
    {
      ROS_WARN("RRTStarGlobalPlanner: failed to find path.");
      return false;
    }

    ros::Time stamp = ros::Time::now();

    for (size_t i = 0; i < path_points.size(); ++i)
    {
      geometry_msgs::PoseStamped pose;
      pose.header.stamp = stamp;
      pose.header.frame_id = global_frame_;
      pose.pose.position.x = path_points[i].x;
      pose.pose.position.y = path_points[i].y;
      pose.pose.position.z = 0.0;
      pose.pose.orientation = goal.pose.orientation;
      plan.push_back(pose);
    }

    nav_msgs::Path path_msg;
    path_msg.header.stamp = stamp;
    path_msg.header.frame_id = global_frame_;
    path_msg.poses = plan;
    plan_pub_.publish(path_msg);

    ROS_INFO("RRTStarGlobalPlanner: generated path with %lu poses.", plan.size());
    return true;
  }

private:
  bool planRRTStar(const Point2D& start,
                   const Point2D& goal,
                   std::vector<Point2D>& final_path)
  {
    std::vector<TreeNode> tree;
    tree.push_back(TreeNode{start, -1, 0.0});

    std::vector<int> goal_candidates;

    int best_goal_idx = -1;
    double best_goal_cost = std::numeric_limits<double>::infinity();

    int success_count_in_period = 0;

    for (int iter = 1; iter <= max_iter_; ++iter)
    {
      if (iter % 10 == 0)
      {
        success_count_in_period = 0;
      }

      double alpha;
      if (best_goal_idx < 0)
      {
        alpha = 1.0 / (1.0 + std::exp(-(static_cast<double>(success_count_in_period) - 5.0)));
      }
      else
      {
        alpha = 0.0;
      }

      Point2D sample = sampleHybrid(start, goal, alpha);

      if (!isPointFree(sample))
      {
        continue;
      }

      int nearest_idx = nearestNode(tree, sample);
      Point2D nearest_point = tree[nearest_idx].point;
      Point2D new_point = steerTowards(nearest_point, sample, step_size_);

      if (!isPointFree(new_point))
      {
        continue;
      }

      if (edgeInCollision(nearest_point, new_point))
      {
        continue;
      }

      success_count_in_period++;

      std::vector<int> neighbor_idx = nearNodes(tree, new_point, search_radius_);

      int reachest_idx = findReachest(tree, nearest_idx, new_point);

      bool has_create = false;
      Point2D x_create;
      int x_create_idx = -1;

      createNode(tree, reachest_idx, new_point, x_create, has_create);

      if (has_create)
      {
        int parent_idx = tree[reachest_idx].parent;

        if (parent_idx >= 0 && !edgeInCollision(tree[parent_idx].point, x_create))
        {
          double create_cost = tree[parent_idx].cost + distance(tree[parent_idx].point, x_create);
          tree.push_back(TreeNode{x_create, parent_idx, create_cost});
          x_create_idx = static_cast<int>(tree.size()) - 1;
        }
        else
        {
          has_create = false;
        }
      }

      int best_parent = nearest_idx;
      double best_cost = tree[nearest_idx].cost + distance(nearest_point, new_point);

      for (size_t i = 0; i < neighbor_idx.size(); ++i)
      {
        int idx = neighbor_idx[i];
        double candidate_cost = tree[idx].cost + distance(tree[idx].point, new_point);

        if (candidate_cost < best_cost && !edgeInCollision(tree[idx].point, new_point))
        {
          best_parent = idx;
          best_cost = candidate_cost;
        }
      }

      tree.push_back(TreeNode{new_point, best_parent, best_cost});
      int new_idx = static_cast<int>(tree.size()) - 1;

      for (size_t i = 0; i < neighbor_idx.size(); ++i)
      {
        int idx = neighbor_idx[i];

        if (idx == new_idx || idx == best_parent)
        {
          continue;
        }

        bool rewired = false;

        if (has_create && x_create_idx >= 0)
        {
          double create_cost = tree[x_create_idx].cost + distance(tree[x_create_idx].point, tree[idx].point);

          if (create_cost < tree[idx].cost && !edgeInCollision(tree[x_create_idx].point, tree[idx].point))
          {
            tree[idx].parent = x_create_idx;
            tree[idx].cost = create_cost;
            updateDescendantCosts(tree, idx);
            rewired = true;
          }
        }

        if (rewired)
        {
          continue;
        }

        double new_cost = tree[new_idx].cost + distance(tree[new_idx].point, tree[idx].point);

        if (new_cost < tree[idx].cost && !edgeInCollision(tree[new_idx].point, tree[idx].point))
        {
          tree[idx].parent = new_idx;
          tree[idx].cost = new_cost;
          updateDescendantCosts(tree, idx);
        }
      }

      if (distance(new_point, goal) <= goal_radius_ && !edgeInCollision(new_point, goal))
      {
        goal_candidates.push_back(new_idx);
      }

      best_goal_idx = -1;
      best_goal_cost = std::numeric_limits<double>::infinity();

      for (size_t i = 0; i < goal_candidates.size(); ++i)
      {
        int idx = goal_candidates[i];

        if (edgeInCollision(tree[idx].point, goal))
        {
          continue;
        }

        double candidate_cost = tree[idx].cost + distance(tree[idx].point, goal);

        if (candidate_cost < best_goal_cost)
        {
          best_goal_cost = candidate_cost;
          best_goal_idx = idx;
        }
      }

      if (stop_at_first_path_ && best_goal_idx >= 0)
      {
        break;
      }
    }

    if (best_goal_idx < 0)
    {
      return false;
    }

    final_path = backtrackPath(tree, best_goal_idx, goal);
    return true;
  }

  Point2D sampleUniform()
  {
    double origin_x = costmap_->getOriginX();
    double origin_y = costmap_->getOriginY();
    double size_x = costmap_->getSizeInMetersX();
    double size_y = costmap_->getSizeInMetersY();

    std::uniform_real_distribution<double> dist_x(origin_x, origin_x + size_x);
    std::uniform_real_distribution<double> dist_y(origin_y, origin_y + size_y);

    return Point2D{dist_x(rng_), dist_y(rng_)};
  }

  Point2D sampleNormal(const Point2D& start, const Point2D& goal)
  {
    Point2D L{goal.x - start.x, goal.y - start.y};
    double L_length = std::sqrt(L.x * L.x + L.y * L.y);

    if (L_length < 1e-9)
    {
      return start;
    }

    double sigma = normal_width_ratio_ * L_length / 4.0;

    std::uniform_real_distribution<double> uniform_01(0.0, 1.0);
    std::normal_distribution<double> normal_dist(0.0, sigma);

    for (int retry = 0; retry < max_sample_retry_; ++retry)
    {
      double t = uniform_01(rng_);

      Point2D lrand{
        start.x + t * L.x,
        start.y + t * L.y
      };

      Point2D sample{
        lrand.x + normal_dist(rng_),
        lrand.y + normal_dist(rng_)
      };

      if (isInsideMap(sample))
      {
        return sample;
      }
    }

    return sampleUniform();
  }

  Point2D sampleHybrid(const Point2D& start, const Point2D& goal, double alpha)
  {
    std::uniform_real_distribution<double> uniform_01(0.0, 1.0);
    double r = uniform_01(rng_);

    if (r > alpha)
    {
      Point2D qrand1 = sampleNormal(start, goal);
      Point2D qrand2 = sampleUniform();

      if (distance(qrand1, goal) <= distance(qrand2, goal))
      {
        return qrand1;
      }
      else
      {
        return qrand2;
      }
    }
    else
    {
      return goal;
    }
  }

  int nearestNode(const std::vector<TreeNode>& tree, const Point2D& point)
  {
    int best_idx = 0;
    double best_dist = std::numeric_limits<double>::infinity();

    for (size_t i = 0; i < tree.size(); ++i)
    {
      double d = distanceSquared(tree[i].point, point);

      if (d < best_dist)
      {
        best_dist = d;
        best_idx = static_cast<int>(i);
      }
    }

    return best_idx;
  }

  std::vector<int> nearNodes(const std::vector<TreeNode>& tree,
                             const Point2D& point,
                             double radius)
  {
    std::vector<int> result;
    double r2 = radius * radius;

    for (size_t i = 0; i < tree.size(); ++i)
    {
      if (distanceSquared(tree[i].point, point) <= r2)
      {
        result.push_back(static_cast<int>(i));
      }
    }

    return result;
  }

  Point2D steerTowards(const Point2D& from, const Point2D& to, double step_size)
  {
    double dx = to.x - from.x;
    double dy = to.y - from.y;
    double d = std::sqrt(dx * dx + dy * dy);

    if (d <= step_size)
    {
      return to;
    }

    return Point2D{
      from.x + dx / d * step_size,
      from.y + dy / d * step_size
    };
  }

  int findReachest(const std::vector<TreeNode>& tree,
                   int nearest_idx,
                   const Point2D& x_rand)
  {
    int reachest_idx = nearest_idx;

    while (tree[reachest_idx].parent >= 0)
    {
      int parent_idx = tree[reachest_idx].parent;
      Point2D parent_point = tree[parent_idx].point;

      if (edgeInCollision(x_rand, parent_point))
      {
        return reachest_idx;
      }

      reachest_idx = parent_idx;
    }

    return reachest_idx;
  }

  void createNode(const std::vector<TreeNode>& tree,
                  int reachest_idx,
                  const Point2D& x_rand,
                  Point2D& x_create,
                  bool& has_create)
  {
    has_create = false;

    int parent_idx = tree[reachest_idx].parent;

    if (parent_idx < 0)
    {
      return;
    }

    Point2D x_allow = tree[reachest_idx].point;
    Point2D x_forbid = tree[parent_idx].point;

    while (distance(x_allow, x_forbid) > d_dichotomy_)
    {
      Point2D x_mid{
        (x_allow.x + x_forbid.x) / 2.0,
        (x_allow.y + x_forbid.y) / 2.0
      };

      if (!edgeInCollision(x_rand, x_mid))
      {
        x_allow = x_mid;
      }
      else
      {
        x_forbid = x_mid;
      }
    }

    x_forbid = x_rand;
    Point2D parent_point = tree[parent_idx].point;

    while (distance(x_allow, x_forbid) > d_dichotomy_)
    {
      Point2D x_mid{
        (x_allow.x + x_forbid.x) / 2.0,
        (x_allow.y + x_forbid.y) / 2.0
      };

      if (!edgeInCollision(x_mid, parent_point))
      {
        x_allow = x_mid;
      }
      else
      {
        x_forbid = x_mid;
      }
    }

    if (distance(x_allow, tree[reachest_idx].point) > 1e-9)
    {
      x_create = x_allow;
      has_create = true;
    }
  }

  void updateDescendantCosts(std::vector<TreeNode>& tree, int root_idx)
  {
    for (size_t i = 0; i < tree.size(); ++i)
    {
      if (tree[i].parent == root_idx)
      {
        tree[i].cost = tree[root_idx].cost + distance(tree[root_idx].point, tree[i].point);
        updateDescendantCosts(tree, static_cast<int>(i));
      }
    }
  }

  std::vector<Point2D> backtrackPath(const std::vector<TreeNode>& tree,
                                     int goal_parent_idx,
                                     const Point2D& goal)
  {
    std::vector<Point2D> reversed_path;
    reversed_path.push_back(goal);

    int idx = goal_parent_idx;

    while (idx >= 0)
    {
      reversed_path.push_back(tree[idx].point);
      idx = tree[idx].parent;
    }

    std::reverse(reversed_path.begin(), reversed_path.end());
    return reversed_path;
  }

  bool isInsideMap(const Point2D& p)
  {
    unsigned int mx;
    unsigned int my;
    return costmap_->worldToMap(p.x, p.y, mx, my);
  }

  bool isPointFree(const Point2D& p)
  {
    unsigned int mx;
    unsigned int my;

    if (!costmap_->worldToMap(p.x, p.y, mx, my))
    {
      return false;
    }

    unsigned char cost = costmap_->getCost(mx, my);

    if (cost == costmap_2d::LETHAL_OBSTACLE)
    {
      return false;
    }

    if (cost == costmap_2d::INSCRIBED_INFLATED_OBSTACLE)
    {
      return false;
    }

    if (!allow_unknown_ && cost == costmap_2d::NO_INFORMATION)
    {
      return false;
    }

    return true;
  }

  bool edgeInCollision(const Point2D& p1, const Point2D& p2)
  {
    double d = distance(p1, p2);

    if (d < 1e-9)
    {
      return !isPointFree(p1);
    }

    int steps = std::max(2, static_cast<int>(std::ceil(d / collision_check_step_)));

    for (int i = 0; i <= steps; ++i)
    {
      double t = static_cast<double>(i) / static_cast<double>(steps);

      Point2D p{
        p1.x + t * (p2.x - p1.x),
        p1.y + t * (p2.y - p1.y)
      };

      if (!isPointFree(p))
      {
        return true;
      }
    }

    return false;
  }

  double distance(const Point2D& a, const Point2D& b)
  {
    return std::sqrt(distanceSquared(a, b));
  }

  double distanceSquared(const Point2D& a, const Point2D& b)
  {
    double dx = a.x - b.x;
    double dy = a.y - b.y;
    return dx * dx + dy * dy;
  }

private:
  bool initialized_;

  costmap_2d::Costmap2DROS* costmap_ros_;
  costmap_2d::Costmap2D* costmap_;
  std::string global_frame_;
  ros::Publisher plan_pub_;

  int max_iter_;
  double step_size_;
  double search_radius_;
  double goal_radius_;
  double normal_width_ratio_;
  int max_sample_retry_;
  double d_dichotomy_;
  bool stop_at_first_path_;
  bool allow_unknown_;
  double collision_check_step_;

  std::mt19937 rng_;
};

}  // namespace rrt_star_global_planner

PLUGINLIB_EXPORT_CLASS(rrt_star_global_planner::RRTStarGlobalPlanner,
                       nav_core::BaseGlobalPlanner)